Update app.py

app.py

@@ -38,8 +38,6 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
     ignore_mismatched_sizes=True)
 
-eos_token_id = tokenizer.eos_token_id
-
 @spaces.GPU()
 def stream_chat(
     message: str,
@@ -64,20 +62,25 @@ def stream_chat(
 
     conversation.append({"role": "user", "content": message})
 
-    input_text = tokenizer.apply_chat_template(conversation, tokenize=False)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    input_text = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(device)
+    
+    terminators = [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("")
+    ]
+
     streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
-        input_ids=inputs,
+        input_ids=input_text,
         max_new_tokens=max_new_tokens,
+        eos_token_id=terminators,
         do_sample=False if temperature == 0 else True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
         streamer=streamer,
-        pad_token_id=eos_token_id,
-        eos_token_id=eos_token_id,
+        pad_token_id=10,
     )
 
     with torch.no_grad():
@@ -88,8 +91,6 @@ def stream_chat(
     for new_text in streamer:
         buffer += new_text
         yield buffer
-        if eos_token_id in tokenizer.encode(new_text):
-            break
 
 
 chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)