Update app.py

app.py

@@ -1,12 +1,11 @@
 import os
-import threading as Thread
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 
-
+MODEL_LIST = ["internlm/internlm2_5-7b-chat", "internlm/internlm2_5-7b-chat-1m"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = os.environ.get("MODEL_ID", None)
 MODEL_NAME = MODEL_ID.split("/")[-1]
@@ -44,40 +43,21 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = model.eval()
 
 @spaces.GPU()
-def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):   
-    conversation = []
-    for prompt, answer in history:
-        conversation.extend([
-            {"role": "user", "content": prompt}, 
-            {"role": "assistant", "content": answer},
-        ])
-    conversation.append({"role": "user", "content": message})
-
-    print(f"Conversation is -\n{conversation}")
-
-    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
+def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
+    print(history)
     
-    streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,})
-
-    generate_kwargs = dict(
-        input_ids=input_ids, 
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        top_k=top_k,
-        repetition_penalty=penalty,
-        do_sample=True, 
-        temperature=temperature,
-        eos_token_id = [2,92542],
+    resp, hist = model.stream_chat(
+        tokenizer,
+        query = message,
+        history = history,
+        max_new_tokens = max_new_tokens,
+        do_sample = True if temperature == 0 else False
+        top_p = top_p,
+        top_k = top_k,
+        temperature = temperature,
     )
-    
-    thread = Thread(target=model.generate, kwargs=generate_kwargs)
-    thread.start()
 
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
+    yield resp
 
 
 chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)