Update README.md

README.md

@@ -9,6 +9,11 @@ Xueguang Ma, Liang Wang, Nan Yang, Furu Wei, Jimmy Lin, arXiv 2023
 
 This model is fine-tuned from LLaMA-2-13B 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
@@ -30,13 +35,13 @@ def get_model(peft_model_name):
 tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-13b-hf')
 model = get_model('castorini/rankllama-v1-13b-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():