Update README.md

README.md

@@ -3,25 +3,25 @@ license: llama2
 ---
 
 
-# RepLLaMA-7B-Passage
+# RankLLaMA-7B-Passage
 
 [Fine-Tuning LLaMA for Multi-Stage Text Retrieval](TODO).
 Xueguang Ma, Liang Wang, Nan Yang, Furu Wei, Jimmy Lin, arXiv 2023
 
-This model is fine-tuned from LLaMA-2-7B using LoRA and the embedding size is 4096.
+This model is fine-tuned from LLaMA-2-7B using LoRA for passage reranking.
 
 ## Usage
 
-Below is an example to encode a query and a document, and then compute their similarity using their embedding.
+Below is an example to compute the similarity score of a query-passage pair
 
 ```python
 import torch
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 from peft import PeftModel, PeftConfig
 
 def get_model(peft_model_name):
     config = PeftConfig.from_pretrained(peft_model_name)
-    base_model = AutoModel.from_pretrained(config.base_model_name_or_path)
+    base_model = AutoModelForSequenceClassification.from_pretrained(config.base_model_name_or_path)
     model = PeftModel.from_pretrained(base_model, peft_model_name)
     model = model.merge_and_unload()
     model.eval()
@@ -29,29 +29,21 @@ def get_model(peft_model_name):
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b-hf')
-model = get_model('castorini/repllama-v1-7b-lora-passage')
+model = get_model('castorini/rankllama-v1-7b-lora-passage')
 
-# Define query and document inputs
+# Define a query-document pair
 query = "What is llama?"
 title = "Llama"
-passage = "The llama is a domesticated South American camelid, widely used as a meat and pack animal by Andean cultures since the pre-Columbian era."
-query_input = tokenizer(f'query: {query}</s>', return_tensors='pt')
-document_input = tokenizer(f'passage: {title} {passage}</s>', return_tensors='pt')
+document = "The llama is a domesticated South American camelid, widely used as a meat and pack animal by Andean cultures since the pre-Columbian era."
 
-# Run the model forward to compute embeddings and query-document similarity score
+# Tokenize the query-document pair
+inputs = tokenizer(f'query: {query}', f'document: {title} {document}</s>', return_tensors='pt')
+
+# Run the model forward
 with torch.no_grad():
-    # compute query embedding
-    query_outputs = model(**query_input)
-    query_embedding = query_outputs.last_hidden_state[0][-1]
-    query_embedding = torch.nn.functional.normalize(query_embedding, p=2, dim=0)
-
-    # compute document embedding
-    document_outputs = model(**document_input)
-    document_embeddings = document_outputs.last_hidden_state[0][-1]
-    document_embeddings = torch.nn.functional.normalize(document_embeddings, p=2, dim=0)
-
-    # compute similarity score
-    score = torch.dot(query_embedding, document_embeddings)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    score = logits[0][0]
     print(score)
 
 ```