Update inference.py

MasterAI-EAM · Jun 4, 2023 · d6961e3 · d6961e3
1 parent c7139f0
commit d6961e3
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/inference.py b/inference.py
@@ -1,5 +1,5 @@
-import time, torch
-from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
+import torch
+from transformers import LlamaTokenizer, LlamaForCausalLM
 import sys 
 
 
@@ -40,6 +40,7 @@ def evaluate(instruction,
         temperature=temperature, 
         top_p=top_p, 
         top_k=top_k,
+        **kwargs
     )
     response = tokenizer.decode(generated_ids[0])
     response = process_response(response)