Update README.md

README.md

@@ -1789,31 +1789,14 @@ print(finetuner.cos_sim(embeddings[0], embeddings[1]))
 Use directly with Huggingface Transformers:
 
 ```python
-import torch
-from transformers import AutoModel, AutoTokenizer
-
-
-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
-    )
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import cos_sim
 
 sentences = ['how is the weather today', 'What is the current weather like today?']
 
-# Load model from HuggingFace Hub
-tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embedding-s-en-v1')
-model = AutoModel.from_pretrained('jinaai/jina-embedding-s-en-v1')
-
-with torch.inference_mode():
-    encoded_input = tokenizer(
-        sentences, padding=True, truncation=True, return_tensors='pt'
-    )
-    model_output = model.encoder(**encoded_input)
-    embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+model = SentenceTransformer('jinaai/jina-embedding-s-en-v1')
+embeddings = model.encode(sentences)
+print(cos_sim(embeddings[0], embeddings[1]))
 ```
 
 ## Fine-tuning