Update inference.py

ecnu-sea · Aug 9, 2024 · f6975a5 · f6975a5
1 parent 3103781
commit f6975a5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/inference/inference.py b/inference/inference.py
@@ -47,6 +47,7 @@ def infer_one(mmd_file_path):
     generated_ids = chat_model.generate(encodes,max_new_tokens=8192,do_sample=True)
     response = tokenizer.batch_decode(generated_ids[: , len_input:])[0]
     return response
+
 def run_review(mmd_file_path):
     infer_modelname = model_name.split('/')[-2]
     infer_save_path = "./" + infer_modelname + '/'
@@ -58,5 +59,5 @@ def run_review(mmd_file_path):
 
 
 if __name__ == "__main__":
-    review = run_review("/path/to")
+    review = run_review("/path/to/xxx.mmd")
     print(review)