Update README.md

README.md

@@ -16,14 +16,21 @@ pip install -U transformers==4.42.4 intel-extension-for-pytorch==2.3.100
 Use the following example below to load the model with the transformers library, tokenize the text, run the model, and apply pooling to the output.
 
 ```
-# example embedding code
 import torch
 from transformers import AutoTokenizer, AutoModel
 import intel_extension_for_pytorch as ipex
 
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]
+    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)
+
 # load model
 tokenizer = AutoTokenizer.from_pretrained('Intel/intel-optimized-model-for-embeddings-v1')
-model = AutoModel.from_pretrained('Intel/intel-optimized-model-for-embeddings-v1', torchscript=True)
+model = AutoModel.from_pretrained('Intel/intel-optimized-model-for-embeddings-v1', 
+                                   torchscript=True)
 model.eval()
 
 # do IPEX optimization
@@ -48,14 +55,8 @@ with torch.no_grad(), torch.cpu.amp.autocast(cache_enabled=False,
     # Call model
     tokenized_text = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
     model_output = model(**tokenized_text)
-
-    # Do mean pooling
-    token_embeddings = model_output[0]
-    attention_mask = tokenized_text['attention_mask']
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    output_sum = torch.sum(token_embeddings * input_mask_expanded, 1)
-    embeddings = output_sum / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-    embeddings = [embeddings[0].tolist()]
+    sentence_embeddings = mean_pooling(model_output,tokenized_text['attention_mask'])
+    embeddings = sentence_embeddings[0].tolist()
 
 # Embeddings output
 print(embeddings)