使用 ModelScope Studio 实现可滚动 ChatBot

news2024/10/2 3:42:37

前面的文章提到过 Gradio 的 Chatbot 不能自动滚动的问题，最近看到了 ModelScope Studio 扩展的一些组件，其中 Chatbot 就完美解决了自动滚动的问题，同时还增加了很多更细化的功能，例如可以设置用户和 AI 的头像。官方文档，可以移步至魔搭的官方空间 https://modelscope.cn/studios/modelscope/modelscope-studio

在这里插入图片描述
本文使用的是 Ollama Qwen 2.5，代码如下，使用 LlamaIndex 和 ModelScope 类库，Ollama 可改为自己使用的模型。

# 安装 modelscope studio
pip install modelscope_studio

from llm import get_local_ollama
from typing import List, Optional, Tuple, Dict
from llama_index.core.llms import ChatMessage, MessageRole
import gradio as gr
import modelscope_studio as mgr
from llama_index.core import PromptTemplate

import os

llm = get_local_ollama("qwen2.5:3b")
History = List[Tuple[str, str]]
Messages = List[ChatMessage]
default_system = 'You are QQ, created by Tom. You are a helpful assistant.'


def modify_system_session(system: str) -> str:
    if system is None or len(system) == 0:
        system = default_system
    return system, system, []


def history_to_messages(history: History, system: str) -> Messages:
    messages = [ChatMessage.from_str(role=MessageRole.SYSTEM, content=system)]
    for h in history:
        messages.append(ChatMessage.from_str(role=MessageRole.USER, content=h[0].text))
        messages.append(ChatMessage.from_str(role=MessageRole.ASSISTANT, content=h[1].text))
    return messages


def messages_to_history(messages: Messages) -> Tuple[str, History]:
    assert messages[0].role == MessageRole.SYSTEM
    system = messages[0].content
    history = []
    for q, r in zip(messages[1::2], messages[2::2]):
        history.append([q.content, r.content])
    return system, history

def process_image (img):
    print (img)
    return img

async def load(text, history, system):
    messages = history_to_messages(history= history, system=system)
    messages.append(ChatMessage.from_str(role=MessageRole.USER.value, content=text))
    
    gen = await llm.astream_chat(messages=messages)
    async for a in gen:
        print(a.message.role)
        system, history = messages_to_history(messages + [ChatMessage.from_str(role=a.message.role, content=a.message.content)])
        yield '', history, system
    
with gr.Blocks () as demo:
    with gr.Row():
        with gr.Column(scale=3):
            system_input = gr.Textbox(value=default_system, label="System Setting")
        with gr.Column(scale=1):
            set_sys_button = gr.Button(value="Set System Prompt", scale=2)
        system_state = gr.Textbox(value=default_system, visible=False)
    user_avatar = os.path.join(os.path.dirname(__file__),"user.jpeg")
    bot_avatar = os.path.join(os.path.dirname(__file__),"bot.jpeg")
    chatbot = mgr.Chatbot(avatar_images=[{"name":"user", "avatar": user_avatar}, {"name":"bot", "avatar":bot_avatar}])
    user_input = gr.Textbox()
        
    with gr.Row():
        button = gr.Button("Submit")
    button.click(load, [user_input, chatbot, system_input], [user_input, chatbot, system_input])
    user_input.submit(load, [user_input, chatbot, system_input], [user_input, chatbot, system_input])
    set_sys_button.click(fn=modify_system_session,
                                    inputs=[system_input],
                                    outputs=[system_state, system_input, chatbot])

    gr.Examples(examples=["讲一个中国的神话故事！", "介绍一下故宫", "使用 Python 某只股票 MA5 > MA10"], inputs=[user_input])

    


demo.launch()