Vades/rubert-tiny2-onnx-optim

import torch from optimum.onnxruntime import ORTModelForFeatureExtraction from transformers import AutoTokenizer import numpy as np

model_path = ".../bert-onnx-optim/" file_name = "model_optimized.onnx"

encoder = ORTModelForFeatureExtraction.from_pretrained( model_path, file_name=file_name ) tokenizer = AutoTokenizer.from_pretrained(model_path)

def encode(self, texts, batch_size=8, normalize=True): """ Fetch embeddings for text :input: Any (str/list): text int: batch size bool: need normilize

    :output:
        list: list fetch embeddings
"""
log.debug(f"Data: {texts}")
if isinstance(texts, str):
    texts = [texts]
all_embeddings = []
for i in range(0, len(texts), batch_size):
    batch = texts[i:i+batch_size]
    inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True)

    with torch.no_grad():
        outputs = self.encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)

    if normalize:
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
    all_embeddings.append(embeddings.cpu().numpy())
return np.vstack(all_embeddings)