利用 Gradio 和 Ollama 构建本地化聊天应用

前言

       在本文中，我们将探讨如何使用 Ollama 提供的大模型 API 和 Gradio 库来创建一个流式对话界面。这个界面允许用户与 AI 模型进行实时对话，并且能够根据用户选择的不同模型生成响应。

一、技术选型

1.1 Gradio

       Gradio 是一个开源的 Python 库，用于快速创建简单的机器学习模型界面。它提供了一个简单易用的 API，可以让我们在几分钟内创建一个交互式的 Web 应用程序。Gradio 支持多种输入和输出类型，包括文本、图像、音频等，这使得我们可以轻松地构建各种类型的机器学习模型界面。在我们的项目中，我们选择了 Gradio 作为聊天应用的界面框架，因为它可以帮助我们快速地创建一个用户友好的聊天界面。

1.2 Ollama

       Ollama是一个专为在本地环境中运行和定制大型语言模型而设计的工具。它提供了一个简单而高效的接口，用于创建、运行和管理这些模型，同时还提供了一个丰富的预构建模型库，可以轻松集成到各种应用程序中。Ollama的目标是使大型语言模型的部署和交互变得简单，无论是对于开发者还是对于终端用户。

二、环境准备

• Python 3.8+
• Ollama
• Gradio 库
• OpenAI库
• requests库

三、代码实现

3.1 获取本地模型列表

       通过调用 Ollama 的 API 接口来获取模型信息。以下是获取本地模型列表的代码：

import requests

def get_local_models():
    api_url = "http://localhost:11434/api/tags"
    try:
        response = requests.get(api_url)
        if response.status_code == 200:
            response_data = response.json()
            models = [model['name'] for model in response_data.get("models", [])]
            return models
        else:
            return []
    except Exception as e:
        print(f"Request failed: {e}")
        return []

3.2 过滤支持聊天的模型

       在获取到本地模型列表后，我们需要过滤出支持聊天功能的模型。我们通过调用 Ollama 的 chat.completions.create 接口来测试每个模型是否支持聊天功能。以下是过滤支持聊天的模型的代码：

from openai import OpenAI

def filter_chat_models(models):
    chat_models = []
    for model in models:
        try:
            client = OpenAI(
                api_key="ollama",
                base_url="http://localhost:11434/v1",
            )
            client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": "你是人工智能助手。"},
                    {
                        "role": "user",
                        "content": "Hello"
                    }
                ],
                stream=False,
            )
            chat_models.append(model)
        except Exception as e:
            print(f"Model {model} does not support chat: {e}")
    return chat_models

3.3 初始化 OpenAI 客户端

       使用 OpenAI 的客户端来调用 Ollama 的 API 接口。

client = OpenAI(
    api_key="ollama",
    base_url="http://localhost:11434/v1",
)

3.4 实现流式回答

       为了提高用户体验，我们在聊天应用中实现了流式回答功能。流式回答可以让我们在生成回答的过程中实时地将内容展示给用户。以下是实现流式回答的代码：

def get_response(text, model_name, history):
    response = client.chat.completions.create(
        model=model_name,
        messages=[
            {"role": "system", "content": "你是人工智能助手。"},
            {
                "role": "user",
                "content": f"{text} " + "以下是历史对话" + f"{history}"
            }
        ],
        stream=True,
    )
    return response

3.5 聊天响应

       在聊天应用中，我们需要实现一个聊天响应函数，用于处理用户的输入并生成回答。以下是聊天响应的代码：

def chat_response(message, history, model_name):
    print(message, history)
    results = get_response(message, model_name, history)
    contents = ""
    for chunk in results:
        if chunk.choices[0].delta.content is not None:
            contents += chunk.choices[0].delta.content
            yield contents

3.6 创建 Gradio 自定义聊天界面

       最后，我们使用 Gradio 创建一个自定义的聊天界面。在聊天界面中，我们提供了一个下拉菜单，让用户可以选择不同的模型。以下是创建 Gradio 聊天界面的代码：

import gradio as gr

with gr.Blocks() as demo:
    chat_interface = gr.ChatInterface(chat_response, title="Your Chat with Ollama", additional_inputs=[gr.Dropdown(label="Model", choices=chat_models, value=models[0] if models else "")])

demo.launch()

四、实际操作

4.1 访问应用界面

       应用启动后，Gradio 会生成一个本地 URL，通常为 http://127.0.0.1:7860。在浏览器中访问该 URL，你将看到聊天应用的界面。

五、总结

       通过上述代码，我们成功地构建了一个基于 Gradio 和 Ollama 的本地化聊天应用。该应用不仅可以实现高效的聊天功能，还可以保护用户的隐私和数据安全。同时，我们还实现了模型选择和过滤功能，让用户可以根据自己的需求选择不同的模型，并确保选择的模型能够正常工作。
       在未来的开发中，我们可以进一步优化该应用，例如增加更多的模型支持、优化界面设计等，以满足更多用户的需求。
       希望本文能够帮助你了解如何使用 Gradio 和 Ollama 构建本地化聊天应用。如果你有任何问题或建议，欢迎随时与我们联系。