Nanbeige4.1-3B WebUI定制化改造指南：添加历史记录/导出功能/多模型切换扩展

Nanbeige4.1-3B WebUI定制化改造指南：添加历史记录/导出功能/多模型切换扩展

如果你已经成功部署了Nanbeige4.1-3B的WebUI，并且用了一段时间，可能会发现一个“痛点”：每次对话都是独立的，聊完就没了，想回顾一下之前的精彩对话，或者想把对话内容保存下来，都得手动复制粘贴，非常麻烦。

更别提，如果你想试试其他模型，还得关掉当前服务，重新配置启动另一个，来回切换很不方便。

这篇文章，就是来解决这些“痛点”的。我将手把手带你，对Nanbeige4.1-3B的官方WebUI进行一番“魔改”，给它加上三个非常实用的功能：

1. 对话历史记录：自动保存你和模型的每一次对话，随时可以查看、继续。
2. 对话内容导出：一键将对话保存为Markdown或TXT文件，方便整理和分享。
3. 多模型热切换：在同一个WebUI界面上，无需重启服务，就能快速切换使用不同的模型。

改造完成后，你的WebUI将从一个基础的对话工具，升级为一个功能更完善、体验更流畅的个人AI工作台。下面，我们就开始动手。

1. 改造前的准备：理解现有代码结构

在动刀之前，我们先得看看“病人”长什么样。根据你提供的项目结构，核心文件是/root/nanbeige-webui/webui.py。

我们先来简单分析一下一个典型的、基于Gradio的LLM WebUI可能的结构（你需要根据你的实际文件内容进行调整）：

# webui.py 可能的核心部分示意 import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer import torch # 1. 模型和分词器加载（通常全局只加载一次） model_path = "/root/ai-models/nanbeige/Nanbeige4___1-3B" tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True ) # 2. 核心的对话生成函数 def predict(message, history, temperature, top_p, max_tokens): """ history 参数：Gradio ChatInterface 会自动传入，格式为 [[用户消息1, AI回复1], [用户消息2, AI回复2], ...] 但基础版可能没利用这个参数做持久化。 """ # 构建对话格式 messages = [] if history: # 如果有历史，将其构建到 messages 中 for human, assistant in history: messages.append({"role": "user", "content": human}) messages.append({"role": "assistant", "content": assistant}) messages.append({"role": "user", "content": message}) # Tokenization 和生成 input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( input_ids, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, do_sample=True ) response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True) # 返回新的回复，Gradio会自动更新界面上的history return response # 3. 创建Gradio界面 with gr.Blocks() as demo: gr.Markdown("# Nanbeige4.1-3B Chat WebUI") chatbot = gr.Chatbot(label="对话历史") # 显示对话的组件 msg = gr.Textbox(label="输入你的问题") # 输入框 with gr.Row(): submit = gr.Button("发送") clear = gr.Button("清空对话") # 参数调节 with gr.Accordion("生成参数", open=False): temperature = gr.Slider(0.0, 2.0, value=0.6, label="Temperature") top_p = gr.Slider(0.0, 1.0, value=0.95, label="Top-P") max_tokens = gr.Slider(128, 4096, value=2048, step=128, label="最大生成长度") # 事件绑定 def respond(message, chat_history, temp, top_p_val, max_tok): bot_message = predict(message, chat_history, temp, top_p_val, max_tok) chat_history.append((message, bot_message)) return "", chat_history # 清空输入框，更新聊天历史 msg.submit(respond, [msg, chatbot, temperature, top_p, max_tokens], [msg, chatbot]) submit.click(respond, [msg, chatbot, temperature, top_p, max_tokens], [msg, chatbot]) clear.click(lambda: None, None, chatbot, queue=False) # 清空界面历史 demo.launch(server_name="0.0.0.0", server_port=7860)

关键点：

• history参数在predict函数和界面交互中流转，但它只存在于内存中，页面一刷新或服务重启就没了。
• 我们的目标是将这个history持久化到本地文件，并增加管理它的界面。

2. 功能一：实现对话历史记录与加载

我们需要一个地方来存储历史对话。一个简单有效的方法是，为每次对话会话创建一个独立的文件，以时间戳或会话ID命名。

2.1 修改webui.py，添加历史记录功能

我们将在文件开头导入必要的库，并添加历史记录管理类。

# webui.py 修改版 - 添加历史记录功能 import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer import torch import json import os from datetime import datetime from typing import List, Tuple, Optional # --- 新增：历史记录管理器 --- class ChatHistoryManager: """管理对话历史记录的保存与加载""" HISTORY_DIR = "./chat_histories" # 历史记录保存目录 def __init__(self): os.makedirs(self.HISTORY_DIR, exist_ok=True) def _get_session_filename(self, session_id: str = None) -> str: """生成会话文件名。如果未提供session_id，则用时间戳生成一个新的。""" if session_id is None: session_id = datetime.now().strftime("session_%Y%m%d_%H%M%S") return os.path.join(self.HISTORY_DIR, f"{session_id}.json") def save_history(self, history: List[Tuple[str, str]], session_id: str = None) -> str: """保存对话历史到JSON文件。返回保存的文件名（session_id）。""" if session_id is None: session_id = datetime.now().strftime("session_%Y%m%d_%H%M%S") filename = self._get_session_filename(session_id) # 将Gradio格式的history [(user, bot), ...] 转换为可序列化的列表 serializable_history = [{"user": h[0], "assistant": h[1]} for h in history] with open(filename, 'w', encoding='utf-8') as f: json.dump({ "session_id": session_id, "created_at": datetime.now().isoformat(), "history": serializable_history }, f, ensure_ascii=False, indent=2) print(f"历史记录已保存至: {filename}") return session_id def load_history(self, session_id: str) -> List[Tuple[str, str]]: """从JSON文件加载对话历史，并转换回Gradio格式。""" filename = self._get_session_filename(session_id) if not os.path.exists(filename): return [] with open(filename, 'r', encoding='utf-8') as f: data = json.load(f) # 转换回 [(user, bot), ...] 格式 history = [(item["user"], item["assistant"]) for item in data["history"]] return history def list_sessions(self) -> List[dict]: """列出所有历史会话。""" sessions = [] for fname in os.listdir(self.HISTORY_DIR): if fname.endswith('.json'): filepath = os.path.join(self.HISTORY_DIR, fname) try: with open(filepath, 'r', encoding='utf-8') as f: data = json.load(f) sessions.append({ "id": data.get("session_id", fname.replace('.json', '')), "created_at": data.get("created_at", ""), "message_count": len(data.get("history", [])) }) except: continue # 按创建时间倒序排列 sessions.sort(key=lambda x: x.get("created_at", ""), reverse=True) return sessions def delete_session(self, session_id: str) -> bool: """删除指定的历史会话。""" filename = self._get_session_filename(session_id) if os.path.exists(filename): os.remove(filename) print(f"已删除会话: {session_id}") return True return False # 初始化历史记录管理器 history_manager = ChatHistoryManager() # --- 原有模型加载部分（暂时不变）--- model_path = "/root/ai-models/nanbeige/Nanbeige4___1-3B" tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True ) # --- 修改预测函数，使其能接收并利用历史 --- def predict_with_history(message, history, temperature, top_p, max_tokens): """支持历史上下文的预测函数""" # 构建messages messages = [] for human, assistant in history: messages.append({"role": "user", "content": human}) messages.append({"role": "assistant", "content": assistant}) messages.append({"role": "user", "content": message}) # Tokenization 和生成（与之前相同） input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( input_ids, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, do_sample=True ) response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True) return response # --- 构建Gradio界面 --- with gr.Blocks(title="Nanbeige4.1-3B 增强版 WebUI", theme=gr.themes.Soft()) as demo: # 使用一个变量来跟踪当前会话ID current_session_id = gr.State(value=None) gr.Markdown(""" # 🚀 Nanbeige4.1-3B 增强版对话界面 **新增功能**：历史记录 | 导出对话 | 多模型切换 """) with gr.Row(): with gr.Column(scale=3): # 主聊天区域 chatbot = gr.Chatbot(label="对话历史", height=500) msg = gr.Textbox(label="输入消息", placeholder="输入您的问题...", lines=2) with gr.Row(): submit_btn = gr.Button("发送", variant="primary") clear_btn = gr.Button("清空当前对话") save_btn = gr.Button("保存当前对话") # 参数调节区域 with gr.Accordion("生成参数设置", open=False): temperature = gr.Slider(0.0, 2.0, value=0.6, label="Temperature (创造性)") top_p = gr.Slider(0.0, 1.0, value=0.95, label="Top-P (多样性)") max_tokens = gr.Slider(128, 8192, value=2048, step=128, label="最大生成长度") with gr.Column(scale=1): # --- 新增：历史记录管理侧边栏 --- with gr.Group(): gr.Markdown("### 💾 历史记录管理") with gr.Row(): session_dropdown = gr.Dropdown( choices=[], label="选择历史会话", interactive=True, info="选择要加载的对话" ) refresh_btn = gr.Button("🔄", size="sm") load_btn = gr.Button("加载选中会话", variant="secondary") delete_btn = gr.Button("删除选中会话", variant="stop") # 会话信息显示 session_info = gr.Markdown("当前会话：新对话") # --- 新增：导出功能区域 --- with gr.Group(): gr.Markdown("### 📤 导出对话") export_format = gr.Radio( choices=["Markdown (.md)", "纯文本 (.txt)", "JSON (.json)"], value="Markdown (.md)", label="导出格式" ) export_btn = gr.Button("导出当前对话", variant="secondary") export_status = gr.Markdown("") # --- 事件处理函数 --- def respond(message, chat_history, session_id, temp, top_p_val, max_tok): """处理用户发送消息""" if not message.strip(): return "", chat_history, session_id bot_message = predict_with_history(message, chat_history, temp, top_p_val, max_tok) chat_history.append((message, bot_message)) # 每次响应后，如果还没有session_id，则生成一个临时的（实际保存时才固定） if session_id is None: session_id = datetime.now().strftime("session_%Y%m%d_%H%M%S") return "", chat_history, session_id def save_current_chat(chat_history, session_id): """保存当前对话到文件""" if not chat_history: return "没有对话内容可保存", session_id # 如果session_id是临时的，通过保存操作获得一个正式的id saved_id = history_manager.save_history(chat_history, session_id) # 更新下拉列表 session_list = history_manager.list_sessions() choices = [s["id"] for s in session_list] return f"对话已保存为: {saved_id}", saved_id, gr.Dropdown.update(choices=choices) def load_selected_session(session_id): """加载选中的历史会话""" if not session_id: return [], "请先选择一个会话", None history = history_manager.load_history(session_id) info_text = f"**已加载会话**: {session_id}\\n消息数: {len(history)}" return history, info_text, session_id def delete_selected_session(session_id): """删除选中的历史会话""" if not session_id: return "请先选择一个会话", gr.Dropdown.update() success = history_manager.delete_session(session_id) if success: # 更新下拉列表 session_list = history_manager.list_sessions() choices = [s["id"] for s in session_list] if session_list else [] return f"已删除会话: {session_id}", gr.Dropdown.update(choices=choices, value=None), [] else: return "删除失败，会话可能不存在", gr.Dropdown.update(), [] def refresh_session_list(): """刷新历史会话列表""" session_list = history_manager.list_sessions() choices = [s["id"] for s in session_list] if session_list else [] return gr.Dropdown.update(choices=choices) def export_chat(chat_history, session_id, format_choice): """导出对话内容""" if not chat_history: return "没有对话内容可导出", "" # 根据格式生成内容 if format_choice == "Markdown (.md)": content = "# 对话记录\\n\\n" for i, (user, assistant) in enumerate(chat_history, 1): content += f"## 第{i}轮\\n" content += f"**用户**: {user}\\n\\n" content += f"**AI助手**: {assistant}\\n\\n---\\n\\n" filename = f"{session_id or 'chat_export'}.md" elif format_choice == "纯文本 (.txt)": content = "对话记录\\n" + "="*30 + "\\n" for i, (user, assistant) in enumerate(chat_history, 1): content += f"[第{i}轮]\\n" content += f"用户: {user}\\n" content += f"AI助手: {assistant}\\n" content += "-"*30 + "\\n" filename = f"{session_id or 'chat_export'}.txt" else: # JSON import json export_data = { "session_id": session_id or "unknown", "export_time": datetime.now().isoformat(), "history": [{"user": u, "assistant": a} for u, a in chat_history] } content = json.dumps(export_data, ensure_ascii=False, indent=2) filename = f"{session_id or 'chat_export'}.json" # 保存文件 export_dir = "./exports" os.makedirs(export_dir, exist_ok=True) filepath = os.path.join(export_dir, filename) with open(filepath, 'w', encoding='utf-8') as f: f.write(content) return f"✅ 对话已导出至: `{filepath}`", "" # --- 绑定事件 --- # 发送消息 msg.submit(respond, [msg, chatbot, current_session_id, temperature, top_p, max_tokens], [msg, chatbot, current_session_id]) submit_btn.click(respond, [msg, chatbot, current_session_id, temperature, top_p, max_tokens], [msg, chatbot, current_session_id]) # 清空对话 clear_btn.click(lambda: ([], None), None, [chatbot, current_session_id]) # 保存对话 save_btn.click(save_current_chat, [chatbot, current_session_id], [export_status, current_session_id, session_dropdown]) # 历史记录管理 refresh_btn.click(refresh_session_list, None, session_dropdown) load_btn.click(load_selected_session, session_dropdown, [chatbot, session_info, current_session_id]) delete_btn.click(delete_selected_session, session_dropdown, [export_status, session_dropdown, chatbot]) # 导出对话 export_btn.click(export_chat, [chatbot, current_session_id, export_format], [export_status, export_status]) # 初始化时刷新会话列表 demo.load(refresh_session_list, None, session_dropdown) # 启动应用 if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

关键改动说明：

1. 新增ChatHistoryManager类：负责所有历史记录的保存、加载、列表和删除操作。数据以JSON格式保存在./chat_histories/目录下。
2. 界面布局重组：使用gr.Row()和gr.Column()将界面分为主聊天区和侧边功能栏。
3. 新增UI组件：
   • session_dropdown：下拉框，显示所有历史会话。
   • refresh_btn/load_btn/delete_btn：历史记录的刷新、加载、删除按钮。
   • export_format/export_btn：导出格式选择和导出按钮。
4. 新增事件处理函数：如save_current_chat,load_selected_session,delete_selected_session,export_chat等，将UI操作与后端逻辑连接起来。
5. 状态管理：使用gr.State()来跟踪current_session_id，确保在对话和保存过程中会话ID的一致性。

2.2 测试历史记录功能

1. 保存修改后的webui.py。
2. 重启WebUI服务（如果之前已运行）：

   cd /root/nanbeige-webui ./stop.sh # 如果提供了停止脚本 # 或者使用 supervisorctl supervisorctl restart nanbeige-webui

3. 访问http://你的服务器IP:7860。
4. 进行几轮对话。
5. 点击“保存当前对话”按钮，侧边栏的“选择历史会话”下拉框应该会更新。
6. 尝试从下拉框选择一个历史会话，点击“加载选中会话”，之前的对话应该会显示在聊天框中。
7. 点击“删除选中会话”可以删除不需要的历史记录。

至此，历史记录功能已经实现。接下来，我们实现第二个功能：多模型切换。

3. 功能二：实现多模型热切换

我们希望在不重启WebUI服务的情况下，动态切换使用不同的语言模型。这需要我们对模型加载部分进行改造，使其支持动态加载和卸载。

3.1 修改webui.py，添加模型管理器

我们将创建一个ModelManager类来管理多个模型，并修改界面以支持切换。

# 在 webui.py 的开头，添加必要的导入 # ... 之前的导入 ... import threading from queue import Queue # --- 新增：模型管理器 --- class ModelManager: """管理多个模型的加载、切换和推理""" def __init__(self): self.current_model_name = None self.current_model = None self.current_tokenizer = None self.model_cache = {} # 缓存已加载的模型 {model_name: (model, tokenizer)} self.lock = threading.Lock() # 防止并发加载/切换 # 预定义的模型配置（你可以在这里添加更多模型） self.model_configs = { "Nanbeige4.1-3B": { "path": "/root/ai-models/nanbeige/Nanbeige4___1-3B", "dtype": torch.bfloat16, "description": "3B参数，擅长推理与对话" }, # 示例：添加另一个模型（请确保路径正确） # "Qwen2.5-3B": { # "path": "/root/ai-models/qwen/Qwen2.5-3B", # "dtype": torch.bfloat16, # "description": "Qwen2.5 3B版本" # } } def load_model(self, model_name: str): """加载指定名称的模型""" with self.lock: if model_name == self.current_model_name and self.current_model is not None: print(f"模型 {model_name} 已是当前模型，无需重新加载。") return True # 检查配置是否存在 if model_name not in self.model_configs: print(f"错误：未找到模型 '{model_name}' 的配置。") return False config = self.model_configs[model_name] model_path = config["path"] dtype = config["dtype"] try: print(f"正在加载模型: {model_name} ...") # 如果模型已在缓存中，直接使用 if model_name in self.model_cache: print(f"从缓存中恢复模型: {model_name}") self.current_model, self.current_tokenizer = self.model_cache[model_name] else: # 加载新的模型和分词器 tokenizer = AutoTokenizer.from_pretrained( model_path, trust_remote_code=True ) model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=dtype, device_map="auto", trust_remote_code=True ) # 放入缓存 self.model_cache[model_name] = (model, tokenizer) self.current_model, self.current_tokenizer = model, tokenizer self.current_model_name = model_name print(f"模型切换成功: {model_name}") return True except Exception as e: print(f"加载模型 {model_name} 失败: {e}") return False def get_current_model_info(self): """获取当前模型信息""" if self.current_model_name: config = self.model_configs.get(self.current_model_name, {}) return { "name": self.current_model_name, "path": config.get("path", "未知"), "description": config.get("description", "无描述") } return {"name": "未加载", "path": "未知", "description": "无"} def unload_model(self, model_name: str): """从缓存中卸载模型以释放显存（可选）""" if model_name in self.model_cache and model_name != self.current_model_name: print(f"正在从缓存中卸载模型: {model_name}") model, tokenizer = self.model_cache.pop(model_name) del model del tokenizer torch.cuda.empty_cache() return True return False # 初始化模型管理器 model_manager = ModelManager() # 默认加载第一个模型 default_model = list(model_manager.model_configs.keys())[0] model_manager.load_model(default_model) # --- 修改预测函数，使用模型管理器 --- def predict_with_history_and_model(message, history, temperature, top_p, max_tokens): """支持历史和动态模型的预测函数""" # 从模型管理器获取当前模型和分词器 model = model_manager.current_model tokenizer = model_manager.current_tokenizer if model is None or tokenizer is None: return "错误：模型未正确加载，请检查模型配置或切换模型。" # 构建messages（与之前相同） messages = [] for human, assistant in history: messages.append({"role": "user", "content": human}) messages.append({"role": "assistant", "content": assistant}) messages.append({"role": "user", "content": message}) # Tokenization 和生成 input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( input_ids, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, do_sample=True ) response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True) return response # --- 修改Gradio界面，添加模型切换组件 --- with gr.Blocks(title="Nanbeige4.1-3B 增强版 WebUI", theme=gr.themes.Soft()) as demo: current_session_id = gr.State(value=None) gr.Markdown(""" # 🚀 多功能AI模型对话平台 **功能**：历史记录 | 导出对话 | **多模型热切换** """) with gr.Row(): with gr.Column(scale=3): # --- 新增：模型切换和信息显示行 --- with gr.Row(): model_dropdown = gr.Dropdown( choices=list(model_manager.model_configs.keys()), value=default_model, label="选择AI模型", interactive=True ) model_switch_btn = gr.Button("切换模型", variant="primary", size="sm") model_info_display = gr.Markdown(f"**当前模型**: {default_model}") # 主聊天区域（保持不变） chatbot = gr.Chatbot(label="对话历史", height=500) msg = gr.Textbox(label="输入消息", placeholder="输入您的问题...", lines=2) with gr.Row(): submit_btn = gr.Button("发送", variant="primary") clear_btn = gr.Button("清空当前对话") save_btn = gr.Button("保存当前对话") # 参数调节区域（保持不变） with gr.Accordion("生成参数设置", open=False): temperature = gr.Slider(0.0, 2.0, value=0.6, label="Temperature (创造性)") top_p = gr.Slider(0.0, 1.0, value=0.95, label="Top-P (多样性)") max_tokens = gr.Slider(128, 8192, value=2048, step=128, label="最大生成长度") with gr.Column(scale=1): # 历史记录管理侧边栏（保持不变） with gr.Group(): gr.Markdown("### 💾 历史记录管理") with gr.Row(): session_dropdown = gr.Dropdown(choices=[], label="选择历史会话", interactive=True) refresh_btn = gr.Button("🔄", size="sm") load_btn = gr.Button("加载选中会话", variant="secondary") delete_btn = gr.Button("删除选中会话", variant="stop") session_info = gr.Markdown("当前会话：新对话") # 导出功能区域（保持不变） with gr.Group(): gr.Markdown("### 📤 导出对话") export_format = gr.Radio( choices=["Markdown (.md)", "纯文本 (.txt)", "JSON (.json)"], value="Markdown (.md)", label="导出格式" ) export_btn = gr.Button("导出当前对话", variant="secondary") export_status = gr.Markdown("") # --- 新增：模型切换事件处理函数 --- def switch_model(model_name): """切换模型""" if not model_name: return "请选择一个模型", model_info_display.update() success = model_manager.load_model(model_name) if success: info = model_manager.get_current_model_info() info_text = f"**当前模型**: {info['name']}\\n**描述**: {info['description']}" return f"✅ 已切换到模型: {model_name}", gr.Markdown.update(value=info_text) else: return f"❌ 切换模型失败: {model_name}", model_info_display.update() # --- 修改响应函数，使用新的预测函数 --- def respond(message, chat_history, session_id, temp, top_p_val, max_tok): if not message.strip(): return "", chat_history, session_id # 使用新的支持模型切换的预测函数 bot_message = predict_with_history_and_model(message, chat_history, temp, top_p_val, max_tok) chat_history.append((message, bot_message)) if session_id is None: session_id = datetime.now().strftime("session_%Y%m%d_%H%M%S") return "", chat_history, session_id # --- 绑定所有事件 --- # 发送消息（使用新的respond函数） msg.submit(respond, [msg, chatbot, current_session_id, temperature, top_p, max_tokens], [msg, chatbot, current_session_id]) submit_btn.click(respond, [msg, chatbot, current_session_id, temperature, top_p, max_tokens], [msg, chatbot, current_session_id]) # 清空对话 clear_btn.click(lambda: ([], None), None, [chatbot, current_session_id]) # 保存对话（保持不变） save_btn.click(save_current_chat, [chatbot, current_session_id], [export_status, current_session_id, session_dropdown]) # 历史记录管理（保持不变） refresh_btn.click(refresh_session_list, None, session_dropdown) load_btn.click(load_selected_session, session_dropdown, [chatbot, session_info, current_session_id]) delete_btn.click(delete_selected_session, session_dropdown, [export_status, session_dropdown, chatbot]) # 导出对话（保持不变） export_btn.click(export_chat, [chatbot, current_session_id, export_format], [export_status, export_status]) # --- 绑定模型切换事件 --- model_switch_btn.click(switch_model, model_dropdown, [export_status, model_info_display]) # 初始化 demo.load(refresh_session_list, None, session_dropdown) # 初始化模型信息显示 initial_info = model_manager.get_current_model_info() initial_info_text = f"**当前模型**: {initial_info['name']}\\n**描述**: {initial_info['description']}" demo.load(lambda: gr.Markdown.update(value=initial_info_text), None, model_info_display) # 启动应用 if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

3.2 配置你自己的模型

关键部分在ModelManager类的model_configs字典。你需要根据你服务器上实际存放的模型路径来修改它。

self.model_configs = { "Nanbeige4.1-3B": { "path": "/root/ai-models/nanbeige/Nanbeige4___1-3B", # 你的模型路径 "dtype": torch.bfloat16, "description": "3B参数，擅长推理与对话" }, "你的第二个模型名称": { "path": "/path/to/your/second/model", # 第二个模型的本地路径 "dtype": torch.bfloat16, # 或 torch.float16 "description": "模型的简要描述" }, # ... 可以继续添加更多模型 }

注意事项：

• 显存占用：同时加载多个大模型会占用大量显存。ModelManager使用了缓存机制，但切换模型时，如果新模型不在缓存中，仍需加载时间并占用额外显存。请根据你的GPU显存量力而行。
• 模型兼容性：确保你添加的模型与transformers库兼容，并且支持apply_chat_template方法（或者你需要调整predict_with_history_and_model函数中的对话构建逻辑）。
• 首次加载：首次切换到一个新模型时，会因为下载/加载模型而有一定延迟。

3.3 测试多模型切换功能

1. 按照上述说明修改model_configs，添加你已有的其他模型路径。
2. 保存webui.py并重启服务。
3. 访问WebUI，你应该能在顶部看到一个“选择AI模型”的下拉框。
4. 从下拉框选择另一个模型，点击“切换模型”按钮。
5. 观察状态提示，如果切换成功，model_info_display区域会更新。
6. 切换后，继续进行对话，模型应该会使用新切换的模型来生成回复。

4. 功能整合与最终优化

现在，我们已经将三个功能整合到了一个WebUI中。为了更好的体验，我们还可以做一些优化：

4.1 添加模型卸载按钮（可选）

在侧边栏增加一个按钮，用于手动释放不用的模型缓存，节省显存。

# 在模型切换行旁边或侧边栏添加 with gr.Row(): model_dropdown = gr.Dropdown(...) model_switch_btn = gr.Button("切换模型", variant="primary", size="sm") model_unload_btn = gr.Button("释放模型缓存", variant="secondary", size="sm") # 新增 model_info_display = gr.Markdown(...) # 添加对应的事件处理函数 def unload_unused_models(): """卸载非当前模型以释放显存""" current = model_manager.current_model_name unloaded = [] for model_name in list(model_manager.model_cache.keys()): if model_name != current: if model_manager.unload_model(model_name): unloaded.append(model_name) if unloaded: return f"已释放模型缓存: {', '.join(unloaded)}" else: return "没有可释放的模型缓存。" # 绑定事件 model_unload_btn.click(unload_unused_models, None, export_status)

4.2 完善错误处理

在实际的预测函数predict_with_history_and_model中，添加更详细的错误处理，避免因为模型加载或生成问题导致整个界面卡死。

def predict_with_history_and_model(message, history, temperature, top_p, max_tokens): """支持历史和动态模型的预测函数（带错误处理）""" model = model_manager.current_model tokenizer = model_manager.current_tokenizer if model is None or tokenizer is None: return "错误：模型未正确加载。请尝试切换模型或检查控制台日志。" try: # 构建messages messages = [] for human, assistant in history: messages.append({"role": "user", "content": human}) messages.append({"role": "assistant", "content": assistant}) messages.append({"role": "user", "content": message}) # Tokenization input_ids = tokenizer.apply_chat_template( messages, return_tensors="pt", add_generation_prompt=True # 确保添加了生成提示 ).to(model.device) # 生成 with torch.no_grad(): outputs = model.generate( input_ids, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, do_sample=True, pad_token_id=tokenizer.eos_token_id # 防止警告 ) # 解码 response = tokenizer.decode( outputs[0][len(input_ids[0]):], skip_special_tokens=True ) return response except torch.cuda.OutOfMemoryError: return "错误：GPU显存不足。请尝试减少生成长度(max_tokens)，或切换/释放其他模型。" except Exception as e: return f"生成过程中出现错误: {str(e)}"

4.3 最终的项目结构

完成所有修改后，你的项目目录可能如下所示：

/root/nanbeige-webui/ ├── webui.py # 改造后的主程序文件 ├── start.sh # 启动脚本 ├── stop.sh # 停止脚本 ├── supervisord.conf # Supervisor配置 ├── requirements.txt # 依赖文件（确保包含 gradio>=4.0） ├── chat_histories/ # 自动创建的，存放历史会话JSON文件 │ ├── session_20250225_143022.json │ └── ... ├── exports/ # 自动创建的，存放导出的对话文件 │ ├── session_20250225_143022.md │ └── ... └── README.md # 可选的更新说明文档

5. 总结与使用建议

通过以上步骤，我们成功地将一个基础的Nanbeige4.1-3B WebUI，改造成为一个具备历史记录、对话导出和多模型热切换三大增强功能的AI对话平台。

5.1 核心改造回顾

1. 历史记录：通过ChatHistoryManager类，将内存中的对话持久化到本地JSON文件，并提供了加载、删除、列表查看的完整界面。
2. 对话导出：在历史记录的基础上，增加了将对话内容一键导出为 Markdown、纯文本或 JSON 格式的功能，便于知识整理和分享。
3. 多模型切换：通过ModelManager类，实现了模型的动态加载和缓存管理。你可以在一个WebUI中轻松切换使用不同的开源大模型，无需重启服务。

5.2 使用建议与后续扩展

• 模型管理：在model_configs字典中仔细配置每个模型的本地路径。建议从参数量小、显存占用低的模型开始测试。
• 显存监控：如果你的GPU显存有限，注意不要同时配置太多大模型。可以善用“释放模型缓存”功能。
• 功能扩展：这个框架具有良好的扩展性。你还可以考虑添加以下功能：
  • 参数预设：保存几组常用的temperature、top_p等参数组合，一键应用。
  • 对话重命名：为历史会话提供自定义名称，而不是只用时间戳。
  • 模型信息页：为每个模型添加更详细的技术说明和示例。
  • 流式输出：将生成过程改为流式（streaming），提升交互体验。

5.3 如何部署与更新

1. 备份：在修改webui.py前，建议先备份原文件。
2. 安装依赖：确保你的环境已安装gradio（版本建议4.0以上）。

   pip install gradio>=4.0.0

3. 重启服务：每次修改webui.py后，需要重启Gradio服务才能生效。

   # 使用你的管理脚本或Supervisor cd /root/nanbeige-webui ./stop.sh ./start.sh # 或 supervisorctl restart nanbeige-webui

4. 查看日志：如果遇到问题，查看日志是首要的排查手段。

   tail -f /var/log/supervisor/nanbeige-webui-stderr.log

希望这份详细的改造指南能帮助你打造一个更加强大、便捷的个人AI对话环境。动手试试吧，享受自定义和扩展开源工具带来的乐趣！

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。