Update README.md

README.md

@@ -10,9 +10,13 @@ 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.
 
+## Training Data
+The model is fine-tuned on the training split of [MS MARCO Passage Ranking](https://microsoft.github.io/msmarco/Datasets) datasets for 1 epoch.
+Please check our paper for details.
+
 ## 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 encode a query and a passage, and then compute their similarity using their embedding.
 
 ```python
 import torch
@@ -31,27 +35,27 @@ def get_model(peft_model_name):
 tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b-hf')
 model = get_model('castorini/repllama-v1-7b-lora-passage')
 
-# Define query and document inputs
+# Define query and passage inputs
 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')
+passage_input = tokenizer(f'passage: {title} {passage}</s>', return_tensors='pt')
 
-# Run the model forward to compute embeddings and query-document similarity score
+# Run the model forward to compute embeddings and query-passage similarity score
 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 passage embedding
+    passage_outputs = model(**passage_input)
+    passage_embeddings = passage_outputs.last_hidden_state[0][-1]
+    passage_embeddings = torch.nn.functional.normalize(passage_embeddings, p=2, dim=0)
 
     # compute similarity score
-    score = torch.dot(query_embedding, document_embeddings)
+    score = torch.dot(query_embedding, passage_embeddings)
     print(score)
 
 ```