doberst commited on
Commit
3584152
1 Parent(s): f6c3974

Update generation_test_hf_script.py

Browse files
Files changed (1) hide show
  1. generation_test_hf_script.py +4 -2
generation_test_hf_script.py CHANGED
@@ -13,7 +13,7 @@ def load_rag_benchmark_tester_ds():
13
 
14
  dataset = load_dataset(ds_name)
15
 
16
- print("update: loading test dataset - ", dataset)
17
 
18
  test_set = []
19
  for i, samples in enumerate(dataset["train"]):
@@ -29,7 +29,9 @@ def run_test(model_name, test_ds):
29
 
30
  device = "cuda" if torch.cuda.is_available() else "cpu"
31
 
32
- print("update: model will be loaded on device - ", device)
 
 
33
 
34
  model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
35
  model.to(device)
 
13
 
14
  dataset = load_dataset(ds_name)
15
 
16
+ print("update: loading RAG Benchmark test dataset - ", dataset)
17
 
18
  test_set = []
19
  for i, samples in enumerate(dataset["train"]):
 
29
 
30
  device = "cuda" if torch.cuda.is_available() else "cpu"
31
 
32
+ print("\nRAG Performance Test - 200 questions")
33
+ print("update: model - ", model_name)
34
+ print("update: device - ", device)
35
 
36
  model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
37
  model.to(device)