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)