Update app.py

app.py

@@ -42,20 +42,31 @@ Your responses should be friendly, unbiased, informative, detailed, and faithful
 system_prompt = f"<|im_start|>{system_role}\n{system_prompt}<|im_end|>"
 
 # Function to generate model predictions.
-
 @spaces.GPU()
 def predict(message, history):
-    # history = []
+    # 初始化对话历史格式
+    if history is None:
+        history = []
+
+    # 在历史中添加当前用户输入，临时设置机器人的回复为空
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()
 
-    # Formatting the input for the model.
-    messages =  system_prompt + sft_end_token.join([sft_end_token.join([f"\n{sft_start_token}{user_role}\n" + item[0], f"\n{sft_start_token}{assistant_role}\n" + item[1]])
-                        for item in history_transformer_format])
+    # 格式化输入为模型需要的格式
+    messages = (
+        system_prompt
+        + sft_end_token.join([
+            sft_end_token.join([
+                f"\n{sft_start_token}{user_role}\n" + item[0],
+                f"\n{sft_start_token}{assistant_role}\n" + item[1]
+            ]) for item in history_transformer_format
+        ])
+    )
     model_inputs = tokenizer([messages], return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        model_inputs,
+        input_ids=model_inputs["input_ids"],
+        attention_mask=model_inputs["attention_mask"],
         streamer=streamer,
         max_new_tokens=1024,
         do_sample=True,
@@ -66,14 +77,23 @@ def predict(message, history):
         stopping_criteria=StoppingCriteriaList([stop]),
         repetition_penalty=1.1,
     )
+
+    # 使用线程来运行生成过程
     t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()  # Starting the generation in a separate thread.
+    t.start()
+
+    # 实时生成部分消息
     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.
+        if sft_end_token in partial_message:  # 检测到停止标志
             break
-        yield partial_message
+        yield history + [[message, partial_message]]  # 输出流式数据
+
+    # 处理生成的最终回复
+    final_message = partial_message.replace(sft_end_token, "").strip()
+    history.append([message, final_message])  # 更新历史记录
+    yield history  # 返回完整对话历史
 
 
 css = """