Update README.md

README.md

@@ -10,6 +10,11 @@ Xueguang Ma, Liang Wang, Nan Yang, Furu Wei, Jimmy Lin, arXiv 2023
 
 This model is fine-tuned from LLaMA-2-7B using LoRA for passage reranking.
 
+## 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 compute the similarity score of a query-passage pair
@@ -31,13 +36,13 @@ def get_model(peft_model_name):
 tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b-hf')
 model = get_model('castorini/rankllama-v1-7b-lora-passage')
 
-# Define a query-document pair
+# Define a query-passage pair
 query = "What is llama?"
 title = "Llama"
-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."
+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."
 
-# Tokenize the query-document pair
-inputs = tokenizer(f'query: {query}', f'document: {title} {document}</s>', return_tensors='pt')
+# Tokenize the query-passage pair
+inputs = tokenizer(f'query: {query}', f'document: {title} {passage}</s>', return_tensors='pt')
 
 # Run the model forward
 with torch.no_grad():