Update generation_test_hf_script.py

generation_test_hf_script.py

@@ -13,7 +13,7 @@ def load_rag_benchmark_tester_ds():
 
     dataset = load_dataset(ds_name)
 
-    print("update: loading test dataset - ", dataset)
+    print("update: loading RAG Benchmark test dataset - ", dataset)
 
     test_set = []
     for i, samples in enumerate(dataset["train"]):
@@ -29,7 +29,9 @@ def run_test(model_name, test_ds):
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    print("update: model will be loaded on device - ", device)
+    print("\nRAG Performance Test - 200 questions")  
+    print("update: model - ", model_name)
+    print("update: device - ", device)
 
     model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
     model.to(device)
@@ -39,6 +41,7 @@ def run_test(model_name, test_ds):
     for i, entries in enumerate(test_ds):
 
         # prepare prompt packaging used in fine-tuning process
+        # note: in out testing, Yi model performed better with trailing "\n" at end of prompt
         new_prompt = "<human>: " + entries["context"] + "\n" + entries["query"] + "\n" + "<bot>:" + "\n"
 
         inputs = tokenizer(new_prompt, return_tensors="pt")