Update app.py

app.py

@@ -1,56 +1,102 @@
 import gradio as gr
 from main import init,clip,answer
-# from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
-# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 
-model,tokenizer = init()
+model,tokenizer = init("attention_lstm_pre.ckpt")
+la_model,la_tokenizer = init("attention_lstm_last.ckpt")
 
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    res = answer(message,model,tokenizer)
-    if res[1]>res[0]:
-        return "unsafe" # unsafe
-    else:
-        return "safe" # safe
-    # messages = [{"role": "system", "content": system_message}]
+# def respond(
+#     message,
+#     history: list[tuple[str, str]],
+#     system_message,
+#     max_tokens,
+#     temperature,
+#     top_p,
+# ):
+#     res = answer(message,model,tokenizer)
+#     if res[1]>res[0]:
+#         return "unsafe" # unsafe
+#     else:
+#         messages = [{"role": "system", "content": system_message}]
+    
+#         for val in history:
+#             if val[0]:
+#                 messages.append({"role": "user", "content": val[0]})
+#             if val[1]:
+#                 messages.append({"role": "assistant", "content": val[1]})
+    
+#         messages.append({"role": "user", "content": message})
+    
+#         response = ""
+    
+#         for message in client.chat_completion(
+#             messages,
+#             max_tokens=max_tokens,
+#             stream=True,
+#             temperature=temperature,
+#             top_p=top_p,
+#         ):
+#             token = message.choices[0].delta.content
+    
+#             response += token
+#             yield response
+#             # 收集所有部分响应
+
 
+def generate_response(messages, max_tokens, temperature, top_p):
+    response = ""
+    
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        response = ""
+        token = message.choices[0].delta.content
+        response += token
+        yield response
+
+def collect_response(message, history, system_message, max_tokens, temperature, top_p):
+    # 创建消息列表
+    messages = [{"role": "system", "content": system_message}]
+    
     # for val in history:
     #     if val[0]:
     #         messages.append({"role": "user", "content": val[0]})
     #     if val[1]:
     #         messages.append({"role": "assistant", "content": val[1]})
+    
+    messages.append({"role": "user", "content": message})
+    # 收集所有部分响应
+    full_response = ""
+    for partial_response in generate_response(messages, max_tokens, temperature, top_p):
+        full_response += partial_response
+    return full_response
+    
 
-    # messages.append({"role": "user", "content": message})
-
-    # response = ""
-
-    # for message in client.chat_completion(
-    #     messages,
-    #     max_tokens=max_tokens,
-    #     stream=True,
-    #     temperature=temperature,
-    #     top_p=top_p,
-    # ):
-    #     token = message.choices[0].delta.content
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    res = answer(message, model, tokenizer)
+    if res[1] > res[0]:
+        return "unsafe"  # unsafe
+    else:
+        # 收集并返回完整的响应
+        full_response = collect_response(message, history, system_message, max_tokens, temperature, top_p)
+        ress = answer(full_response,la_model,la_tokenizer)
+        if res[1] > res[0]:
+            return "unsafe"  # unsafe
+        else:
+            return full_response
 
-    #     response += token
-    #     yield response
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[