Update app.py

eos token set to im_end and debugging

app.py

@@ -25,6 +25,7 @@ class ChatBot:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 MODEL_NAME, trust_remote_code=True
             )
+            print(self.tokenizer.eos_token_id)
             print("Tokenizer loaded!")
 
             print("Loading model...")
@@ -68,11 +69,6 @@ class ChatBot:
         # Tokenize
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
 
-        # Get stop token IDs for "assistant\n"
-        stop_token_ids = self.tokenizer.encode(
-            "assistant\n", add_special_tokens=False
-        )
-
         # Generate response
         with torch.no_grad():
             outputs = self.model.generate(
@@ -82,7 +78,7 @@ class ChatBot:
                 do_sample=True,
                 return_dict_in_generate=True,
                 pad_token_id=self.tokenizer.eos_token_id,
-                eos_token_id=[self.tokenizer.eos_token_id] + stop_token_ids,
+                eos_token_id=[self.tokenizer.eos_token_id] + "<|im_end|>",
             )
 
         # Decode and clean response, with multiple debugs