Update README.md

README.md

@@ -34,49 +34,52 @@ The `gte-modernbert` models demonstrates competitive performance in several text
 ## Usage
 
 Use with `Transformers`
-
 ```python
 # Requires transformers>=4.36.0
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-import torch.nn.functional as F
-from transformers import AutoModel, AutoTokenizer
-
-input_texts = [
-    "what is the capital of China?",
-    "how to implement quick sort in python?",
-    "Beijing",
-    "sorting algorithms"
-]
-
-model_path = 'Alibaba-NLP/gte-modernbert-base'
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
-
-# Tokenize the input texts
-batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')
-
-outputs = model(**batch_dict)
-embeddings = outputs.last_hidden_state[:, 0]
- 
-# (Optionally) normalize embeddings
-embeddings = F.normalize(embeddings, p=2, dim=1)
-scores = (embeddings[:1] @ embeddings[1:].T) * 100
-print(scores.tolist())
-```
+model_name_or_path = 'Alibaba-NLP/gte-reranker-modernbert-base'
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+model = AutoModelForSequenceClassification.from_pretrained(
+    model_name_or_path, trust_remote_code=True,
+    torch_dtype=torch.float16
+)
+model.eval()
+
+pairs = [["what is the capital of China?", "Beijing"], ["how to implement quick sort in python?","Introduction of quick sort"], ["how to implement quick sort in python?", "The weather is nice today"]]
 
+with torch.no_grad():
+    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
+    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
+    print(scores)
+
+# tensor([1.2315, 0.5923, 0.3041])
+```
 Use with `sentence-transformers`:
 
+Before you start, install the sentence-transformers libraries:
+```
+pip install sentence-transformers
+```
+
 ```python
 # Requires sentence_transformers>=2.7.0
+from sentence_transformers import CrossEncoder
+
+model_name_or_path = 'Alibaba-NLP/gte-reranker-modernbert-base'
+
+model = CrossEncoder(
+    model_name_or_path,
+    automodel_args={"torch_dtype": "auto"},
+    trust_remote_code=True,
+)
 
-from sentence_transformers import SentenceTransformer
-from sentence_transformers.util import cos_sim
+pairs = [["what is the capital of China?", "Beijing"], ["how to implement quick sort in python?","Introduction of quick sort"], ["how to implement quick sort in python?", "The weather is nice today"]]
 
-sentences = ['That is a happy person', 'That is a very happy person']
+scores = model.predict(sentence_pairs, convert_to_tensor=True).tolist()
 
-model = SentenceTransformer('Alibaba-NLP/gte-modernbert-base', trust_remote_code=True)
-embeddings = model.encode(sentences)
-print(cos_sim(embeddings[0], embeddings[1]))
+print ("scores: ", scores)
 ```
 
 ## Training Details