Fix for termination issue (#38)

ymcui · May 13, 2024 · 4871fc1 · 4871fc1
1 parent 0ca21bd
commit 4871fc1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/inference/inference_hf.py b/scripts/inference/inference_hf.py
@@ -191,7 +191,7 @@ def generate_prompt(instruction):
                     generation_output = model.generate(
                         input_ids = inputs["input_ids"].to(device),
                         attention_mask = inputs['attention_mask'].to(device),
-                        eos_token_id=tokenizer.eos_token_id,
+                        eos_token_id=terminators,
                         pad_token_id=tokenizer.eos_token_id,
                         generation_config = generation_config
                     )
@@ -215,4 +215,4 @@ def generate_prompt(instruction):
                 with open(dirname+'/generation_config.json','w') as f:
                     json.dump(generation_config,f,ensure_ascii=False,indent=2)
             else:
-                generation_config.save_pretrained('./')
+                generation_config.save_pretrained('./')