Update app.py

app.py

@@ -1,64 +1,225 @@
+import os
+import time
+
 import gradio as gr
-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")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    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})
-
-    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
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+import modelscope_studio.components.antd as antd
+import modelscope_studio.components.antdx as antdx
+import modelscope_studio.components.base as ms
+import modelscope_studio.components.pro as pro
+from modelscope_studio.components.pro.chatbot import (ChatbotBotConfig,
+                                                      ChatbotPromptsConfig,
+                                                      ChatbotUserConfig,
+                                                      ChatbotWelcomeConfig)
+from openai import OpenAI
+
+gw_api_key = os.getenv("GW_API_KEY")
+
+client = OpenAI(
+    base_url='http://api-test.geniuworks.com/v2',
+    api_key=gw_api_key',  # ModelScope Token
 )
 
+model = "xinyuan-32b-v0609"
+
+
+def prompt_select(e: gr.EventData):
+    return gr.update(value=e._data["payload"][0]["value"]["description"])
+
+
+def clear():
+    return gr.update(value=None)
+
+
+def retry(chatbot_value, e: gr.EventData):
+    index = e._data["payload"][0]["index"]
+    chatbot_value = chatbot_value[:index]
+
+    yield gr.update(loading=True), gr.update(value=chatbot_value), gr.update(
+        disabled=True)
+    for chunk in submit(None, chatbot_value):
+        yield chunk
+
+
+def cancel(chatbot_value):
+    chatbot_value[-1]["loading"] = False
+    chatbot_value[-1]["status"] = "done"
+    chatbot_value[-1]["footer"] = "Chat completion paused"
+    return gr.update(value=chatbot_value), gr.update(loading=False), gr.update(
+        disabled=False)
+
+
+def format_history(history):
+    messages = []
+    for item in history:
+        if item["role"] == "user":
+            messages.append({"role": "user", "content": item["content"]})
+        elif item["role"] == "assistant":
+            # ignore thought message
+            messages.append({
+                "role": "assistant",
+                "content": item["content"][-1]["content"]
+            })
+    return messages
+
+
+def submit(sender_value, chatbot_value):
+    if sender_value is not None:
+        chatbot_value.append({
+            "role": "user",
+            "content": sender_value,
+        })
+    history_messages = format_history(chatbot_value)
+    chatbot_value.append({
+        "role": "assistant",
+        "content": [],
+        "loading": True,
+        "status": "pending"
+    })
+    yield {
+        sender: gr.update(value=None, loading=True),
+        clear_btn: gr.update(disabled=True),
+        chatbot: gr.update(value=chatbot_value)
+    }
+
+    try:
+        response = client.chat.completions.create(model=model,
+                                                  messages=history_messages,
+                                                  stream=True,
+                                                  max_tokens=2048,
+                                                  temperature=0.6,
+                                                  top_p=0.95,
+                                                  )
+        thought_done = False
+        start_time = time.time()
+        message_content = chatbot_value[-1]["content"]
+        # thought content
+        message_content.append({
+            "copyable": False,
+            "editable": False,
+            "type": "tool",
+            "content": "",
+            "options": {
+                "title": "Thinking..."
+            }
+        })
+        # content
+        message_content.append({
+            "type": "text",
+            "content": "",
+        })
+        for chunk in response:
+            reasoning_content = chunk.choices[0].delta.reasoning_content
+            content = chunk.choices[0].delta.content
+            chatbot_value[-1]["loading"] = False
+            message_content[-2]["content"] += reasoning_content or ""
+            message_content[-1]["content"] += content or ""
+
+            if content and not thought_done:
+                thought_done = True
+                thought_cost_time = "{:.2f}".format(time.time() - start_time)
+                message_content[-2]["options"][
+                    "title"] = f"End of Thought ({thought_cost_time}s)"
+                message_content[-2]["options"]["status"] = "done"
+
+            yield {chatbot: gr.update(value=chatbot_value)}
+
+        chatbot_value[-1]["footer"] = "{:.2f}".format(time.time() -
+                                                      start_time) + 's'
+        chatbot_value[-1]["status"] = "done"
+        yield {
+            clear_btn: gr.update(disabled=False),
+            sender: gr.update(loading=False),
+            chatbot: gr.update(value=chatbot_value),
+        }
+    except Exception as e:
+        chatbot_value[-1]["loading"] = False
+        chatbot_value[-1]["status"] = "done"
+        chatbot_value[-1]["content"] = "Failed to respond, please try again."
+        yield {
+            clear_btn: gr.update(disabled=False),
+            sender: gr.update(loading=False),
+            chatbot: gr.update(value=chatbot_value),
+        }
+        raise e
+
+
+with gr.Blocks() as demo, ms.Application(), antdx.XProvider():
+    with antd.Flex(vertical=True, gap="middle"):
+        chatbot = pro.Chatbot(
+            height=800,
+            welcome_config=ChatbotWelcomeConfig(
+                variant="borderless",
+                icon=
+                "./xinyuan.png",
+                title=f"Hello, I'm Xinyuan👋",
+                description="You can input text to get started.",
+                prompts=ChatbotPromptsConfig(
+                    title="How can I help you today?",
+                    styles={
+                        "list": {
+                            "width": '100%',
+                        },
+                        "item": {
+                            "flex": 1,
+                        },
+                    },
+                    items=[{
+                        "label":
+                        "💝 心理学与实际应用",
+                        "children": [{
+                            "description":
+                            "课题分离是什么意思？"
+                        }, {
+                            "description":
+                            "回避型依恋和焦虑型依恋有什么区别？还有其他依恋类型吗？"
+                        }, {
+                            "description":
+                            "为什么我背单词的时候总是只记得开头和结尾，中间全忘了？"
+                        }]
+                    }, {
+                        "label":
+                        "👪 儿童教育与发展",
+                        "children": [{
+                            "description":
+                            "什么是正念养育？"
+                        }, {
+                            "description":
+                            "2岁孩子分离焦虑严重，送托育中心天天哭闹怎么办？"
+                        }, {
+                            "description":
+                            "4岁娃说话不清还爱打人，是心理问题还是欠管教？"
+                        }]
+                    }])),
+            user_config=ChatbotUserConfig(
+                avatar="https://api.dicebear.com/7.x/miniavs/svg?seed=3",
+                variant="shadow"),
+            bot_config=ChatbotBotConfig(
+                header=model,
+                avatar=
+                "./xinyuan.png",
+                actions=["copy", "retry"],
+                variant="shadow"),
+        )
+
+        with antdx.Sender() as sender:
+            with ms.Slot("prefix"):
+                with antd.Button(value=None, color="default",
+                                 variant="text") as clear_btn:
+                    with ms.Slot("icon"):
+                        antd.Icon("ClearOutlined")
+        clear_btn.click(fn=clear, outputs=[chatbot])
+        submit_event = sender.submit(fn=submit,
+                                     inputs=[sender, chatbot],
+                                     outputs=[sender, chatbot, clear_btn])
+        sender.cancel(fn=cancel,
+                      inputs=[chatbot],
+                      outputs=[chatbot, sender, clear_btn],
+                      cancels=[submit_event],
+                      queue=False)
+        chatbot.retry(fn=retry,
+                      inputs=[chatbot],
+                      outputs=[sender, chatbot, clear_btn])
+        chatbot.welcome_prompt_select(fn=prompt_select, outputs=[sender])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()