File size: 1,676 Bytes
2349468 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
from typing import Dict, List, Any
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer
import torch.nn.functional as F
import torch
# copied from the model card
def mean_pooling(model_output, attention_mask):
token_embeddings = model_output[0] #First element of model_output contains all token embeddings
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
class sentence_embeddings(path = '.'):
def __init__(self, path):
# load the optimized model
self.model = ORTModelForFeatureExtraction.from_pretrained(path, file_name="model_quantized.onnx")
self.tokenizer = AutoTokenizer.from_pretrained(path)
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
"""
Args:
data (:obj:):
includes the input data and the parameters for the inference.
Return:
A :obj:`list`:. The list contains the embeddings of the inference inputs
"""
inputs = data.get("inputs", data)
# tokenize the input
encoded_inputs = self.tokenizer(inputs, padding=True, truncation=True, return_tensors='pt')
# run the model
outputs = self.model(**encoded_inputs)
# Perform pooling
embeddings = mean_pooling(outputs, encoded_inputs['attention_mask'])
# Normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
# postprocess the prediction
return {'embeddings': embeddings.tolist()}
|