Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteriaList, StoppingCriteria
 import os
 from threading import Thread
 
@@ -34,6 +34,15 @@ model = AutoModelForCausalLM.from_pretrained(
 
 tokenizer = AutoTokenizer.from_pretrained("THUDM/LongWriter-glm4-9b",trust_remote_code=True)
 
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        # stop_ids = model.config.eos_token_id
+        stop_ids = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
+                    tokenizer.get_command("<|observation|>")]
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
 
 @spaces.GPU()
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):
@@ -49,24 +58,29 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
     input_ids = tokenizer.build_chat_input(message, history=conversation, role='user').input_ids.to(model.device)
     #input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
     streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
+                    tokenizer.get_command("<|observation|>")]
 
     generate_kwargs = dict(
+        input_ids=input_ids,
         max_new_tokens=max_new_tokens,
         streamer=streamer,
         do_sample=True,
         top_k=1,
         temperature=temperature,
         repetition_penalty=1,
+        stopping_criteria=StoppingCriteriaList([stop]),
+        eos_token_id=eos_token_id,
     )
-    gen_kwargs = {**input_ids, **generate_kwargs}
+    #gen_kwargs = {**input_ids, **generate_kwargs}
 
-    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
     buffer = ""
     for new_text in streamer:
         buffer += new_text
         yield buffer
- 
+
 chatbot = gr.Chatbot(height=600, placeholder = PLACEHOLDER)
 
 with gr.Blocks(css=CSS) as demo: