Update README.md

README.md

@@ -1,74 +1,73 @@
 ---
+pipeline_tag: sentence-similarity
+language: en
 license: apache-2.0
+tags:
+- sentence-transformers
+- feature-extraction
+- sentence-similarity
+- transformers
+- onnx
 ---
 
+# 
 
-```
-##How to download the model?
-!sudo apt-get install git-lfs
-!git lfs install
+This is the ONNX model of sentence-transformers/all-roberta-large-v1 [https://seb.sbert.net]. Currently, Hugging Face does not support downloading ONNX files with external format files. I have created a workaround using sbert and optimum together to generate embeddings.
 
-# Then
-!git clone https://huggingface.co/vamsibanda/sbert-onnx-all-roberta-large-v1
 ```
-
-## How 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
 ```
-from onnxruntime import InferenceSession
+
+Then you can use the model like this:
+
+```python
+import os
+from torch.hub import _get_torch_home
+from sentence_transformers.util import snapshot_download
+from transformers import AutoTokenizer, pipeline
+from optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTModelForFeatureExtraction
+from sentence_transformers.models import Transformer, Pooling, Dense
 import torch
 from transformers.modeling_outputs import BaseModelOutput
-from transformers import RobertaTokenizerFast
+from transformers import T5TokenizerFast
 import torch.nn.functional as F
-from sentence_transformers.models import Transformer, Pooling, Dense
-
-class RobertaEncoder(torch.nn.Module):
-    def __init__(self, encoder_sess):
-        super().__init__()
-        self.encoder = encoder_sess
 
-    def forward(
-        self,
-        input_ids,
-        attention_mask,
-        inputs_embeds=None,
-        head_mask=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
-    ):
-
-        encoder_hidden_state = torch.from_numpy(
-            self.encoder.run(
-                None,
-                {
-                    "input_ids": input_ids.cpu().numpy(),
-                    "attention_mask": attention_mask.cpu().numpy(),
-                    
-                },
-            )[0]
-        )
-
-        return BaseModelOutput(encoder_hidden_state)
+model_name = 'vamsibanda/sbert-onnx-all-roberta-large-v1'
+cache_folder = './'
+model_path =  os.path.join(cache_folder, model_name.replace("/", "_"))
 
 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 sbert_onnx_encode(sentence_input):
-    token = roberta_tokenizer(sentence_input, return_tensors='pt')
-    encoder_outputs = encoder_layer(input_ids=token['input_ids'], attention_mask=token['attention_mask'])
-    sbert_embeddings = mean_pooling(encoder_outputs, token['attention_mask'])
-    sbert_embeddings = F.normalize(sbert_embeddings, p=2, dim=1)
-    return sbert_embeddings.tolist()[0]
-
-roberta_tokenizer = RobertaTokenizerFast.from_pretrained('sbert-onnx-all-roberta-large-v1')
-encoder_sess = InferenceSession('sbert-onnx-all-roberta-large-v1/sbert-roberta-large-quant.onnx')
-encoder_layer = RobertaEncoder(encoder_sess)
-pooling_layer = Pooling.load('./sbert-onnx-all-roberta-large-v1/1_Pooling/')
+def download_onnx_model(model_name, cache_folder, model_path, force_download = False):
+   if force_download and os.path.exists(model_path):
+      os.remove(model_path)
+  
+   snapshot_download(model_name,
+                     cache_dir=cache_folder,
+                     library_name='sentence-transformers',
+                     ignore_files=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'],
+                    )
+   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")
 
-m1 = sbert_onnx_encode('That is a happy person')
-m2 = sbert.encode('That is a happy person').tolist()
-print(util.cos_sim(m1,m2))
-##tensor([[0.9925]])
+token = tokenizer('That is a happy person', 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)
+sbert_embeddings.tolist()[0]
+                    
 ```