From f37ecda996382a2b3e9eb0dfeaea0a1fdc96ae2c Mon Sep 17 00:00:00 2001 From: Cody Date: Thu, 14 Dec 2023 19:37:33 -0500 Subject: [PATCH] feat: stop phi-2 output at <|endoftext|> --- phi2/phi2.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/phi2/phi2.py b/phi2/phi2.py index 7973c33d..2da78378 100644 --- a/phi2/phi2.py +++ b/phi2/phi2.py @@ -200,8 +200,12 @@ if __name__ == "__main__": print("[INFO] Generating with Phi-2...", flush=True) print(args.prompt, end="", flush=True) + end_of_response = tokenizer.encode("<|endoftext|>")[0] + tokens = [] for token, _ in zip(generate(prompt, model), range(args.max_tokens)): + if token.item() == end_of_response: + break tokens.append(token) if (len(tokens) % 10) == 0: