Update README.md

README.md

@@ -10,65 +10,22 @@ tags:
 - onnx
 ---
 
-# 
-
-This is the ONNX model of sentence-transformers/all-MiniLM-L6-v2 [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.
+# ONNX convert all-MiniLM-L6-v2
+## Conversion of [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2)
+This is a [sentence-transformers](https://www.SBERT.net) ONNX model: It maps sentences & paragraphs to a 384 dimensional dense vector space and can be used for tasks like clustering or semantic search. This custom model takes `last_hidden_state` and `pooler_output` whereas the sentence-transformers exported with default ONNX config only contains `last_hidden_state` as output.
 
+## Usage (HuggingFace Optimum)
+Using this model becomes easy when you have [optimum](https://github.com/huggingface/optimum) installed:
 ```
-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
+python -m pip install optimum
 ```
-
 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-MiniLM-L6-v2'
-cache_folder = './'
-model_path =  os.path.join(cache_folder, model_name.replace("/", "_"))
-
-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
-      
-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 generate_embedding(text):
-    token = tokenizer(text, return_tensors='pt')
-    embedding = model(input_ids=token['input_ids'], attention_mask=token['attention_mask'])
-    embedding = mean_pooling(embedding, token['attention_mask'])
-    embedding = F.normalize(embedding, p=2, dim=1)
-    return embedding.tolist()[0]
-                    
-_ = 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")
+from optimum.onnxruntime.modeling_ort import ORTModelForCustomTasks
 
-generate_embedding('That is a happy person')
-                    
+model = ORTModelForCustomTasks.from_pretrained("vamsibanda/sbert-all-MiniLM-L6-with-pooler")
+tokenizer = AutoTokenizer.from_pretrained("vamsibanda/sbert-all-MiniLM-L6-with-pooler")
+inputs = tokenizer("I love burritos!", return_tensors="pt")
+pred = model(**inputs)
+embedding = pred['pooler_output']
 ```