File size: 2,675 Bytes
63700b0
c2c33f1
 
63700b0
c2c33f1
 
 
 
 
 
63700b0
a286e00
c2c33f1
a286e00
10d6ae9
f645a17
 
c2c33f1
 
 
 
 
 
 
b1e07be
c2c33f1
 
 
 
 
 
b8b4869
 
c2c33f1
a286e00
 
 
eb644f5
a286e00
c2c33f1
 
 
a286e00
8a2d2a7
 
 
 
 
 
 
a286e00
 
 
 
 
c2c33f1
2fc3370
 
 
b6595d7
b8b4869
c2c33f1
c07c56a
b8b4869
 
c2c33f1
 
 
 
 
 
a286e00
8a2d2a7
c2c33f1
b1e07be
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
---
pipeline_tag: sentence-similarity
language: en
license: apache-2.0
tags:
- sentence-transformers
- feature-extraction
- sentence-similarity
- transformers
- onnx
---

# 

This is the ONNX model of sentence-transformers/all-roberta-large-v1 [https://seb.sbert.net]. Currently, Hugging Face does not support downloading ONNX model and generate embeddings. I have created a workaround using sbert and optimum together to generate embeddings.

```
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.0
```

Then you can use the model like this:

```python
import os
from sentence_transformers.util import snapshot_download
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForFeatureExtraction
from sentence_transformers.models import Transformer, Pooling, Dense
import torch
from transformers.modeling_outputs import BaseModelOutput
import torch.nn.functional as F
import shutil

model_name = 'vamsibanda/sbert-onnx-all-roberta-large-v1'
cache_folder = './'
model_path =  os.path.join(cache_folder, model_name.replace("/", "_"))

def generate_embedding(text):
    token = tokenizer(text, return_tensors='pt')
    embeddings = model(input_ids=token['input_ids'], attention_mask=token['attention_mask'])
    sbert_embeddings = mean_pooling(embeddings, token['attention_mask'])
    sbert_embeddings = F.normalize(sbert_embeddings, p=2, dim=1)
    return sbert_embeddings.tolist()[0]
    
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)

def download_onnx_model(model_name, cache_folder, model_path, force_download = False):
    if force_download and os.path.exists(model_path):
       shutil.rmtree(model_path)
    elif os.path.exists(model_path):
       return 
    snapshot_download(model_name,
                     cache_dir=cache_folder,
                     library_name='sentence-transformers'
                     )
    return
                    
     
_ = download_onnx_model(model_name, cache_folder, model_path)                
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = ORTModelForFeatureExtraction.from_pretrained(model_path, force_download=False)
pooling_layer = Pooling.load(f"{model_path}/1_Pooling")

generate_embedding('That is a happy person')
                    
```