From 8d496ba61a54b6f79c301facf4ae0b5096c7d8eb Mon Sep 17 00:00:00 2001
From: devonthomas35 <30363743+devonthomas35@users.noreply.github.com>
Date: Thu, 14 Dec 2023 15:50:59 -0800
Subject: [PATCH] Stop generating at eos token

---
 phi2/phi2.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/phi2/phi2.py b/phi2/phi2.py
index 7973c33d..ede79ea2 100644
--- a/phi2/phi2.py
+++ b/phi2/phi2.py
@@ -202,7 +202,11 @@ if __name__ == "__main__":
 
     tokens = []
     for token, _ in zip(generate(prompt, model), range(args.max_tokens)):
-        tokens.append(token)
+        
+        if token == tokenizer.eos_token_id:
+            break
+        else:
+            tokens.append(token)
 
         if (len(tokens) % 10) == 0:
             mx.eval(tokens)