fix: cutoff prompt correctly

2 years ago · f51c5c8109
parent b6e3ba07c4
commit f51c5c8109
1 changed files with 1 additions and 1 deletions
--- a/generate.py
+++ b/generate.py
@ -11,7 +11,7 @@ def generate(tokenizer, prompt, model, config):

    outputs = model.generate(input_ids=input_ids, max_new_tokens=config["max_new_tokens"], temperature=config["temperature"])

-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    return decoded[len(prompt):]