Update app.py

app.py

@@ -38,6 +38,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
     ignore_mismatched_sizes=True)
 
+eos_token_id = tokenizer.eos_token_id
 
 @spaces.GPU()
 def stream_chat(
@@ -75,7 +76,8 @@ def stream_chat(
         top_k=top_k,
         temperature=temperature,
         streamer=streamer,
-        pad_token_id=10,
+        pad_token_id=eos_token_id,
+        eos_token_id=eos_token_id,
     )
 
     with torch.no_grad():
@@ -86,6 +88,8 @@ 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)