Update README.md

README.md

@@ -8,7 +8,7 @@ pipeline_tag: sentence-similarity
 library_name: transformers
 ---
 
-# gte-modernbert-base
+# gte-reranker-modernbert-base
 
 We are excited to introduce the `gte-modernbert` series of models, which are built upon the latest modernBERT pre-trained encoder-only foundation models. The `gte-modernbert` series models include both text embedding models and rerank models.
 
@@ -21,7 +21,6 @@ The `gte-modernbert` models demonstrates competitive performance in several text
 - Primary Language: English
 - Model Size: 149M
 - Max Input Length: 8192 tokens
-- Output Dimension: 768
 
 ### Model list
 |                                         Models                                         | Language |       Model Type       | Model Size | Max Seq. Length | Dimension | MTEB-en | BEIR | LoCo | CoIR |
@@ -36,71 +35,52 @@ Use with `Transformers`
 ```python
 # Requires transformers>=4.48.0
 
-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())
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+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
 
-from sentence_transformers import SentenceTransformer
-from sentence_transformers.util import cos_sim
+model_name_or_path = 'Alibaba-NLP/gte-reranker-modernbert-base'
 
-sentences = ['That is a happy person', 'That is a very happy person']
+model = CrossEncoder(
+    model_name_or_path,
+    automodel_args={"torch_dtype": "auto"},
+    trust_remote_code=True,
+)
 
-model = SentenceTransformer('Alibaba-NLP/gte-modernbert-base', trust_remote_code=True)
-embeddings = model.encode(sentences)
-print(cos_sim(embeddings[0], embeddings[1]))
-```
+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"]]
+
+scores = model.predict(sentence_pairs, convert_to_tensor=True).tolist()
 
-Use with `transformers.js`:
-
-```js
-// npm i @xenova/transformers
-import { pipeline, dot } from '@xenova/transformers';
-
-// Create feature extraction pipeline
-const extractor = await pipeline('feature-extraction', 'Alibaba-NLP/gte-modernbert-base', {
-    quantized: false, // Comment out this line to use the quantized version
-});
-
-// Generate sentence embeddings
-const sentences = [
-    "what is the capital of China?",
-    "how to implement quick sort in python?",
-    "Beijing",
-    "sorting algorithms"
-]
-const output = await extractor(sentences, { normalize: true, pooling: 'cls' });
-
-// Compute similarity scores
-const [source_embeddings, ...document_embeddings ] = output.tolist();
-const similarities = document_embeddings.map(x => 100 * dot(source_embeddings, x));
-console.log(similarities);
+print ("scores: ", scores)
 ```
 
 ## Training Details