Update README.md

README.md

@@ -13,7 +13,7 @@ datasets:
 
 # Model Summary
 
-Promptriever is a new way of using dense retriever models. This version, `promptriever-llama2-7b-v1` was instruction-trained on a corpus of 490k MSMarco samples with instructions and 490k without instructions. See the [paper]() for more details.
+Promptriever is a bi-encoder retrieval model that can take in natural language instructions and prompts. This version, `promptriever-llama2-7b-v1` was instruction-trained on a corpus of 490k MSMarco samples with instructions and 490k without instructions. See the [paper](todo) for more details.
 
 - **Repository:** [orionw/Promptriever](https://github.com/orionw/promptriever)
 - **Paper:** [Promptriever: Instruction-Trained Retrievers Can Be Prompted Like Language Models](TODO) 
@@ -22,7 +22,16 @@ Promptriever is a new way of using dense retriever models. This version, `prompt
 
 # Use
 
-Below is an example to compute the similarity score of a query-document pair
+You can use MTEB to load this model ([source code](https://github.com/embeddings-benchmark/mteb/blob/main/mteb/models/promptriever_models.py)):
+```python
+import mteb
+model = mteb.get_model("samaya-ai/promptriever-llama2-7b-v1")
+tasks = mteb.get_tasks(tasks=["NFCorpus"], languages=["eng"])
+evaluation = mteb.MTEB(tasks=tasks)
+evaluation.run(model, batch_size=16)
+```
+
+If you want to use a different framework, here's an example of how to batch:
 ```python
 import torch
 import torch.nn.functional as F
@@ -33,7 +42,7 @@ import numpy as np
 class Promptriever:
     def __init__(self, model_name_or_path):
         self.model, self.tokenizer = self.get_model(model_name_or_path)
-        self.model.eval()
+        self.model.eval().cuda()
 
     def get_model(self, peft_model_name):
         # Load the PEFT configuration to get the base model name
@@ -44,33 +53,82 @@ class Promptriever:
         base_model = AutoModel.from_pretrained(base_model_name)
         tokenizer = AutoTokenizer.from_pretrained(base_model_name)
         tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        tokenizer.padding_side = "right"
 
         # Load and merge the PEFT model
         model = PeftModel.from_pretrained(base_model, peft_model_name)
         model = model.merge_and_unload()
 
+        # can be much longer, but for the example 512 is enough
+        model.config.max_length = 512
+        tokenizer.model_max_length = 512
+
         return model, tokenizer
 
-    def encode(self, texts):
-        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        embeddings = outputs.last_hidden_state[:, 0]  # Using [CLS] token
-        return F.normalize(embeddings, p=2, dim=1)
+    def create_batch_dict(self, tokenizer, input_texts):
+        max_length = self.model.config.max_length
+        batch_dict = tokenizer(
+            input_texts,
+            max_length=max_length - 1,
+            return_token_type_ids=False,
+            return_attention_mask=False,
+            padding=False,
+            truncation=True,
+        )
+        batch_dict["input_ids"] = [
+            input_ids + [tokenizer.eos_token_id]
+            for input_ids in batch_dict["input_ids"]
+        ]
+        return tokenizer.pad(
+            batch_dict,
+            padding=True,
+            pad_to_multiple_of=8,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+
+    def encode(self, sentences, max_length: int = 2048, batch_size: int = 4):
+        all_embeddings = []
+        for i in range(0, len(sentences), batch_size):
+            batch_texts = sentences[i : i + batch_size]
+
+            batch_dict = self.create_batch_dict(self.tokenizer, batch_texts)
+            batch_dict = {
+                key: value.to(self.model.device) for key, value in batch_dict.items()
+            }
+
+            with torch.cuda.amp.autocast():
+                with torch.no_grad():
+                    outputs = self.model(**batch_dict)
+                    last_hidden_state = outputs.last_hidden_state
+                    sequence_lengths = batch_dict["attention_mask"].sum(dim=1) - 1
+                    batch_size = last_hidden_state.shape[0]
+                    reps = last_hidden_state[
+                        torch.arange(batch_size, device=last_hidden_state.device),
+                        sequence_lengths,
+                    ]
+                    embeddings = F.normalize(reps, p=2, dim=-1)
+                    all_embeddings.append(embeddings.cpu().numpy())
+
+        return np.concatenate(all_embeddings, axis=0)
 
 # Initialize the model
 model = Promptriever("samaya-ai/promptriever-llama2-7b-v1")
 
 # Example query and instruction
 query = "What universities are in Baltimore, Maryland?"
-instruction = "A relevant document would describe any university in Baltimore. I am only interested in the United States, so ignore any document with a campus in Italy."
 
-# Combine query and instruction with two spaces after "query: "
-input_text = f"query:  {query} {instruction}"
+# add specific relevance conditions if desired (and/or/not) and any other prompts
+instruction = "A relevant document would describe any university in Baltimore. I am not interested in any university that was the first American university. Think carefully about these conditions when determining relevance."
+
+# Combine query and instruction with **two spaces** after "query: "
+input_text = f"query:  {query.strip()} {instruction.strip()}".strip()
 
 # Example documents
-doc1 = "Johns Hopkins University (often abbreviated as Johns Hopkins, Hopkins, or JHU) is a private research university in Baltimore, Maryland. Founded in 1876, Johns Hopkins was the first American university based on the European research institution model. The university also has graduate campuses in Italy, China, and Washington, D.C."
-doc2 = "Johns Hopkins University (often abbreviated as Johns Hopkins, Hopkins, or JHU) is a private research university in Baltimore, Maryland. Founded in 1876, Johns Hopkins was the first American university based on the European research institution model. The university also has graduate campuses in China, and Washington, D.C."
+# NOTE: double space after `passage:`
+doc1 = "passage:  Johns Hopkins University (often abbreviated as Johns Hopkins, Hopkins, or JHU) is a private research university in Baltimore, Maryland. Founded in 1876, Johns Hopkins was the first American university based on the European research institution model."
+doc2 = "passage:  Johns Hopkins University (often abbreviated as Johns Hopkins, Hopkins, or JHU) is a private research university in Baltimore, Maryland. Founded in 1876, Johns Hopkins was the second American university based on the European research institution model."
 
 # Encode query and documents
 query_embedding = model.encode([input_text])
@@ -78,11 +136,17 @@ doc_embeddings = model.encode([doc1, doc2])
 
 # Calculate similarities
 similarities = np.dot(query_embedding, doc_embeddings.T)[0]
+print(f"Similarities: {similarities}") # Similarities: [0.53341305 0.53451955]
+assert similarities[1] > similarities[0]
+
 
-# Print results
-print("Similarities:")
-print(f"Document 1: {similarities[0]:.4f}")
-print(f"Document 2: {similarities[1]:.4f}")
+# change up the instruction to the opposite, to see it works
+instruction = "A relevant document would describe any university in Baltimore. I am interested in any university that was the first American university. Think carefully about these conditions when determining relevance."
+input_text = f"query:  {query.strip()} {instruction.strip()}".strip()
+query_embedding = model.encode([input_text])
+similarities = np.dot(query_embedding, doc_embeddings.T)[0]
+print(f"Similarities: {similarities}") # Similarities: [0.60182875 0.5874183 ]
+assert similarities[0] > similarities[1]
 ```
 
 # Training
@@ -124,7 +188,7 @@ deepspeed --include localhost:0,1,2,3 --master_port "60002" --module tevatron.re
 ```
 
 # License
-This model is released under the Apache-2 license, following the terms of service of the Llama licence. This model was used for research efforts and is not used in any production systems at Samaya AI.  
+This model was used for research efforts and is not used in any production systems at Samaya AI.  Usage must follow the license of the base model as well, as this is a LoRA fine-tune.
 
 # Citation