Upload generation_test_hf_script.py

generation_test_hf_script.py

@@ -27,10 +27,14 @@ def load_rag_benchmark_tester_ds():
 
 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)
+
     model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model.to(device)
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
     for i, entries in enumerate(test_ds):
 
@@ -63,7 +67,7 @@ def run_test(model_name, test_ds):
         bot = output_only.find("<bot>:")
         if bot > -1:
             output_only = output_only[bot+len("<bot>:"):]
-            
+
         # end - post-processing
 
         print("\n")
@@ -78,7 +82,6 @@ if __name__ == "__main__":
     test_ds = load_rag_benchmark_tester_ds()
 
     model_name = "llmware/dragon-red-pajama-7b-v0"
-
     output = run_test(model_name,test_ds)