Update app.py

app.py

@@ -4,8 +4,8 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
 
-# model_path = 'dreamerdeo/Sailor2-0.8B-Chat'
-model_path = 'sail/Sailor-0.5B-Chat'
+model_path = 'dreamerdeo/Sailor2-0.8B-Chat'
+# model_path = 'sail/Sailor-0.5B-Chat'
 
 # Loading the tokenizer and model from Hugging Face's model hub.
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
@@ -78,33 +78,15 @@ def predict(message, history):
         repetition_penalty=1.1,
     )
 
-    outputs = model.generate(**generate_kwargs)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    partial_message = generated_text
-    final_message = partial_message.replace(sft_end_token, "").strip()
-
-    return final_message
-
-    # # 使用线程来运行生成过程
-    # t = Thread(target=model.generate, kwargs=generate_kwargs)
-    # t.start()
-
-    # # 实时生成部分消息
-    # partial_message = ""
-    # for new_token in streamer:
-    #     partial_message += new_token
-    #     if sft_end_token in partial_message:  # 检测到停止标志
-    #         break
-    #     # 将历史记录和当前消息转换为 tuple 格式并实时返回
-    #     # yield [(msg, bot) for msg, bot in history] + [(message, partial_message)]
-    #     # yield (message, partial_message)
-    #     yield partial_message
-    
-    # # 处理生成的最终回复
-    # final_message = partial_message.replace(sft_end_token, "").strip()
-    # history.append([message, final_message])  # 更新历史记录
-    # # 返回最终的对话历史，确保格式为元组的列表
-    # yield [(msg, bot) for msg, bot in history]
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()  # Starting the generation in a separate thread.
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        if sft_end_token in partial_message:  # Breaking the loop if the stop token is generated.
+            break
+        yield partial_message
+
 
 css = """
 full-height {