Views
No views yet
pip install onnx
pip install onnxruntime==1.10.0
pip install transformers>4.6.1
pip install sentencepiece
pip install sentence-transformers
pip install optimum
pip install torch==1.9.01
2import os
3from sentence_transformers.util import snapshot_download
4from transformers import AutoTokenizer
5from optimum.onnxruntime import ORTModelForFeatureExtraction
6from sentence_transformers.models import Transformer, Pooling, Dense
7import torch
8from transformers.modeling_outputs import BaseModelOutput
9import torch.nn.functional as F
10import shutil
11
12model_name = 'vamsibanda/sbert-onnx-gtr-t5-xl'
13cache_folder = './'
14model_path = os.path.join(cache_folder, model_name.replace("/", "_"))
15
16def generate_embedding(text):
17 token = tokenizer(text, return_tensors='pt')
18 embeddings = model(input_ids=token['input_ids'], attention_mask=token['attention_mask'])
19 sbert_embeddings = mean_pooling(embeddings, token['attention_mask'])
20 sbert_embeddings = dense_layer.forward({'sentence_embedding':sbert_embeddings})
21 sbert_embeddings = F.normalize(sbert_embeddings['sentence_embedding'], p=2, dim=1)
22 return sbert_embeddings.tolist()[0]
23
24def download_onnx_model(model_name, cache_folder, model_path, force_download = False):
25 if force_download and os.path.exists(model_path):
26 shutil.rmtree(model_path)
27 elif os.path.exists(model_path):
28 return
29 snapshot_download(model_name,
30 cache_dir=cache_folder,
31 library_name='sentence-transformers'
32 )
33 return
34
35def mean_pooling(model_output, attention_mask):
36 token_embeddings = model_output[0] #First element of model_output contains all token embeddings
37 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
38 return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
39
40def generate_embedding(text):
41 token = tokenizer(text, return_tensors='pt')
42 embedding = model(input_ids=token['input_ids'], attention_mask=token['attention_mask'])
43 embedding = mean_pooling(embedding, token['attention_mask'])
44 embedding = dense_layer.forward({'sentence_embedding':embedding})
45 embedding = F.normalize(embedding['sentence_embedding'], p=2, dim=1)
46 return embedding.tolist()[0]
47
48
49_ = download_onnx_model(model_name, cache_folder, model_path)
50tokenizer = AutoTokenizer.from_pretrained(model_path)
51model = ORTModelForFeatureExtraction.from_pretrained(model_path, force_download=False)
52pooling_layer = Pooling.load(f"{model_path}/1_Pooling")
53dense_layer = Dense.load(f"{model_path}/2_Dense")
54
55generate_embedding('That is a happy person')
56