基于 Gradio 手撸 ollama 模型聊天界面

由于最近在测试各种开源模型性能，在终端测试下载下来的ollama 模型交互较为不方便，搬运工基于Gradio 手撸一个可以启动各种大模型的聊天界面；聊天交互界实现效果如下：

现将实战代码分享如下：

# -*- coding: utf-8 -*-
import gradio as gr
import ollama
from typing import List, Tuple

model_name_dict = {
    "qwen2.5:32b": "qwen2.5:32b",
    "qwen2.5:72b": "qwen2.5:72b",
    "qwen2.5-coder": "qwen2.5-coder:32b",
    "qwq-32b": "qwq:32b-preview-q8_0",
}


# LLMAsServer 模型即服务， 启动聊天界面
class LLMAsServer:
    def __init__(self, model_name: str = model_name_dict["qwen2.5:32b"], is_stream=False):
        self.model_name = model_name
        self.is_stream = is_stream

    def generate(self, text: str):
        stream = ollama.generate(
            stream=self.is_stream,
            model=self.model_name,
            prompt=text,
        )
        response = ""
        if self.is_stream:
            for chunk in stream:
                response += chunk["response"]
        else:
            response = stream["response"]
        return response

    # api_generate 流式输出
    def _api_generate(self, chat_history: List):
        text = chat_history[-1]["content"]
        response = self.generate(text)
        yield chat_history + [{"role": "assistant", "content": response}]

    @staticmethod
    def _handle_user_message(user_message: str, history: list):
        """Handle the user submitted message. Clear message box and append to the history."""
        new_history = history + [{"role": "user", "content": user_message}]
        return '', new_history

    @staticmethod
    def _reset_chat() -> tuple:
        """Reset the agent's chat history. And clear all dialogue boxes."""
        return "", []

    def run(self):
        custom_css = """
            #box {
                height: 500px;
                overflow-y: scroll !important;
            }
            #clear-button {
                background-color: blue !important;
                color: white !important; /* Optional: Set text color to white for better contrast */
                border: none !important; /* Optional: Remove default border */
                border-radius: 2px !important; /* Optional: Add rounded corners */
                height: 150px !important; /* Set the height of the Textbox */
            }
            #message-input {
                width: 100% !important; /* Make the Textbox take up the full available width */
                max-width: calc(100% - 50px) !important; /* Adjust for the ClearButton's width */
                height: 150px !important; /* Increase the height to 50px */
            }
            #submit-button {
                background-color: green !important;
            }
            textarea{
                height: 100px !important; /* Ensure the inner container takes up the full height */
                border-radius: 5px !important; /* Optional: Add rounded corners */
                border-color: red !important; 
                border-style: solid !important; 
            }
            
        """
        demo = gr.Blocks(
            theme=gr.themes.Default(
                primary_hue="red",
                secondary_hue="pink"),
            # css="#box { height: 420px; overflow-y: scroll !important}",
            css=custom_css,
        )
        with demo:
            gr.Markdown(
                "# AI Chat \n"
                "### This Application  is Powered by The LLM Model {} \n".format(self.model_name.upper()),
            )
            chat_window = gr.Chatbot(
                label="Message History",
                scale=5,
                type='messages',
                elem_id="chatbot-container"
            )
            with gr.Row():
                message_input = gr.Textbox(label="Input", placeholder="Type your message here...", scale=5,
                                           elem_id="message-input")
                clear = gr.ClearButton(elem_id="clear-button")
            submit_button = gr.Button("Submit", elem_id="submit-button")

            submit_button.click(
                fn=self._handle_user_message,
                inputs=[message_input, chat_window],
                outputs=[message_input, chat_window]
            ).then(
                fn=self._api_generate,
                inputs=chat_window,
                outputs=chat_window
            )

            message_input.submit(
                self._handle_user_message,
                inputs=[message_input, chat_window],
                outputs=[message_input, chat_window],
            ).then(
                self._api_generate,
                chat_window,
                chat_window,
            )
            clear.click(self._reset_chat, None, [message_input, chat_window])

        demo.launch(show_error=True, share=False, server_port=7861)


if __name__ == "__main__":
    server = LLMAsServer(model_name=model_name_dict["qwen2.5-coder"])
    server.run()

Gradio 构建聊天界面功能强大，能支持CSS 设置，还能编写js 代码插入其中；给构建你的机器学习应用与分享带来了方便。

关注机器学习搬运工，不定期分享最前沿的AI 算法。谢谢大家！