import random import re import json import os import sys from threading import Thread import torch import numpy as np import streamlit as st from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from trainer.trainer_utils import safe_math_eval st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed") st.markdown(""" """, unsafe_allow_html=True) device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") # 多语言文本 LANG_TEXTS = { 'zh': { 'settings': '模型设定调整', 'history_rounds': '历史对话轮次', 'max_length': '最大生成长度', 'temperature': '温度', 'thinking': '思考', 'tools': '工具', 'language': '语言', 'send': '给 MiniMind 发送消息', 'disclaimer': 'AI 生成内容可能存在错误，请仔细核实', 'think_tip': '自适应思考，目前多轮对话或Tool Call共存时思考不稳定', 'tool_select': '工具选择（最多4个）', }, 'en': { 'settings': 'Model Settings', 'history_rounds': 'History Rounds', 'max_length': 'Max Length', 'temperature': 'Temperature', 'thinking': 'Thinking', 'tools': 'Tools', 'language': 'Language', 'send': 'Send a message to MiniMind', 'disclaimer': 'AI-generated content may be inaccurate, please verify', 'think_tip': 'Adaptive thinking; may be unstable with multi-turn or Tool Call', 'tool_select': 'Tool Selection (max 4)', } } def get_text(key): lang = st.session_state.get('lang', 'en') return LANG_TEXTS.get(lang, {}).get(key, LANG_TEXTS['zh'].get(key, key)) # 工具定义 TOOLS = [ {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式", "parameters": {"type": "object", "properties": {"expression": {"type": "string", "description": "数学表达式"}}, "required": ["expression"]}}}, {"type": "function", "function": {"name": "get_current_time", "description": "获取当前时间", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "default": "Asia/Shanghai"}}, "required": []}}}, {"type": "function", "function": {"name": "random_number", "description": "生成随机数", "parameters": {"type": "object", "properties": {"min": {"type": "integer"}, "max": {"type": "integer"}}, "required": ["min", "max"]}}}, {"type": "function", "function": {"name": "text_length", "description": "计算文本长度", "parameters": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}}}, {"type": "function", "function": {"name": "unit_converter", "description": "单位转换", "parameters": {"type": "object", "properties": {"value": {"type": "number"}, "from_unit": {"type": "string"}, "to_unit": {"type": "string"}}, "required": ["value", "from_unit", "to_unit"]}}}, {"type": "function", "function": {"name": "get_current_weather", "description": "获取天气", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}, {"type": "function", "function": {"name": "get_exchange_rate", "description": "获取汇率", "parameters": {"type": "object", "properties": {"from_currency": {"type": "string"}, "to_currency": {"type": "string"}}, "required": ["from_currency", "to_currency"]}}}, {"type": "function", "function": {"name": "translate_text", "description": "翻译文本", "parameters": {"type": "object", "properties": {"text": {"type": "string"}, "target_lang": {"type": "string"}}, "required": ["text", "target_lang"]}}}, ] TOOL_SHORT_NAMES = { 'calculate_math': '数学', 'get_current_time': '时间', 'random_number': '随机', 'text_length': '字数', 'unit_converter': '单位', 'get_current_weather': '天气', 'get_exchange_rate': '汇率', 'translate_text': '翻译' } def execute_tool(tool_name, args): import datetime try: if tool_name == 'calculate_math': return {"result": safe_math_eval(args.get('expression', '0'))} elif tool_name == 'get_current_time': tz = args.get('timezone', 'Asia/Shanghai') return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} elif tool_name == 'random_number': return {"result": random.randint(args.get('min', 0), args.get('max', 100))} elif tool_name == 'text_length': return {"result": len(args.get('text', ''))} elif tool_name == 'unit_converter': return {"result": f"{args.get('value', 0)} {args.get('from_unit', '')} = ? {args.get('to_unit', '')}"} elif tool_name == 'get_current_weather': return {"result": f"{args.get('city', 'Unknown')}: 晴, 7~10°C"} elif tool_name == 'get_exchange_rate': return {"result": f"1 {args.get('from_currency', 'USD')} = 7.2 {args.get('to_currency', 'CNY')}"} elif tool_name == 'translate_text': return {"result": f"[翻译结果]: hello world"} return {"result": "Unknown tool"} except Exception as e: return {"error": str(e)} def process_assistant_content(content, is_streaming=False): # 处理tool_call标签，格式化显示 if '' in content: def format_tool_call(match): try: tc = json.loads(match.group(1)) name = tc.get('name', 'unknown') args = tc.get('arguments', {}) return f'
ToolCalling
{name}: {json.dumps(args, ensure_ascii=False)}
' except: return match.group(0) content = re.sub(r'(.*?)', format_tool_call, content, flags=re.DOTALL) # 流式生成且开启思考时，一开始就放到折叠里 if is_streaming and st.session_state.get('enable_thinking', False) and '' not in content and '' not in content: m = re.search(r'(\n\n(?:我是|您好|你好)[^\n]*)', content) if m and m.start(1) > 5: i = m.start(1) think_part = content[:i] answer_part = content[i:] return f'
已思考
{think_part.strip()}
{answer_part}' elif len(content) > 5: return f'
思考中...
{content.strip().replace(chr(10), "
")}
' if '' in content and '' in content: def format_think(match): think_content = match.group(2) if think_content.replace('\n', '').strip(): # 不是全换行 return f'
已思考
{think_content.strip()}
' return '' content = re.sub(r'()(.*?)()', format_think, content, flags=re.DOTALL) if '' in content and '' not in content: def format_think_in_progress(match): tc = match.group(1) return f'
思考中...
{tc.strip().replace(chr(10), "
")}
' content = re.sub(r'(.*?)$', format_think_in_progress, content, flags=re.DOTALL) if '' not in content and '' in content: def format_think_no_start(match): think_content = match.group(1) if think_content.replace('\n', '').strip(): return f'
已思考
{think_content.strip()}
' return '' content = re.sub(r'(.*?)', format_think_no_start, content, flags=re.DOTALL) return content @st.cache_resource def load_model_tokenizer(model_path): model = AutoModelForCausalLM.from_pretrained( model_path, trust_remote_code=True ) tokenizer = AutoTokenizer.from_pretrained( model_path, trust_remote_code=True ) model = model.half().eval().to(device) return model, tokenizer def clear_chat_messages(): del st.session_state.messages del st.session_state.chat_messages def init_chat_messages(): if "messages" in st.session_state: for i, message in enumerate(st.session_state.messages): if message["role"] == "assistant": st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True) else: st.markdown( f'
{message["content"]}
', unsafe_allow_html=True) else: st.session_state.messages = [] st.session_state.chat_messages = [] return st.session_state.messages def regenerate_answer(index): st.session_state.messages.pop() st.session_state.chat_messages.pop() st.rerun() # 动态扫描模型目录 script_dir = os.path.dirname(os.path.abspath(__file__)) MODEL_PATHS = {} for d in sorted(os.listdir(script_dir), reverse=True): full_path = os.path.join(script_dir, d) if os.path.isdir(full_path) and not d.startswith('.') and not d.startswith('_'): if any(f.endswith(('.bin', '.safetensors', '.pt')) or os.path.exists(os.path.join(full_path, 'model.safetensors.index.json')) for f in os.listdir(full_path) if os.path.isfile(os.path.join(full_path, f))): MODEL_PATHS[d] = [d, d] if not MODEL_PATHS: MODEL_PATHS = {"No models found": ["", "No models"]} # 模型选择 selected_model = st.sidebar.selectbox('Model', list(MODEL_PATHS.keys()), index=0) model_path = MODEL_PATHS[selected_model][0] slogan = f"我是 {MODEL_PATHS[selected_model][1]}，有什么可以帮你的？" if st.session_state.get('lang', 'en') == 'zh' else f"I am {MODEL_PATHS[selected_model][1]}, how can I help you?" st.sidebar.markdown('
', unsafe_allow_html=True) # 语言选择 lang_options = {'中文': 'zh', 'English': 'en'} current_lang = st.session_state.get('lang', 'en') lang_index = 0 if current_lang == 'zh' else 1 lang_label = st.sidebar.radio('Language / 语言', list(lang_options.keys()), index=lang_index, horizontal=True) if lang_options[lang_label] != current_lang: st.session_state.lang = lang_options[lang_label] st.rerun() st.sidebar.markdown('
', unsafe_allow_html=True) # 参数设置 st.session_state.history_chat_num = st.sidebar.slider(get_text('history_rounds'), 0, 8, 0, step=2) st.session_state.max_new_tokens = st.sidebar.slider(get_text('max_length'), 256, 8192, 8192, step=1) st.session_state.temperature = st.sidebar.slider(get_text('temperature'), 0.6, 1.2, 0.90, step=0.01) st.sidebar.markdown('
', unsafe_allow_html=True) # 功能开关 st.session_state.enable_thinking = st.sidebar.checkbox(get_text('thinking'), value=False, help=get_text('think_tip')) st.session_state.selected_tools = [] with st.sidebar.expander(get_text('tools')): st.caption(get_text('tool_select')) selected_count = sum(1 for tool in TOOLS if st.session_state.get(f"tool_{tool['function']['name']}", False)) for tool in TOOLS: name = tool['function']['name'] short_name = TOOL_SHORT_NAMES.get(name, name) checked = st.checkbox(short_name, key=f"tool_{name}", disabled=(selected_count >= 4 and not st.session_state.get(f"tool_{name}", False))) if checked and len(st.session_state.selected_tools) < 4: st.session_state.selected_tools.append(name) image_url = "https://www.modelscope.cn/api/v1/studio/gongjy/MiniMind/repo?Revision=master&FilePath=images%2Flogo2.png&View=true" st.markdown( f'
' '
' f' ' f'{slogan}' '
' f'{get_text("disclaimer")}' '
', unsafe_allow_html=True ) def setup_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def main(): model, tokenizer = load_model_tokenizer(model_path) if "messages" not in st.session_state: st.session_state.messages = [] st.session_state.chat_messages = [] messages = st.session_state.messages for i, message in enumerate(messages): if message["role"] == "assistant": st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True) else: st.markdown( f'
{message["content"]}
', unsafe_allow_html=True) prompt = st.chat_input(key="input", placeholder=get_text('send')) if hasattr(st.session_state, 'regenerate') and st.session_state.regenerate: prompt = st.session_state.last_user_message regenerate_index = st.session_state.regenerate_index delattr(st.session_state, 'regenerate') delattr(st.session_state, 'last_user_message') delattr(st.session_state, 'regenerate_index') if prompt: st.markdown( f'
{prompt}
', unsafe_allow_html=True) messages.append({"role": "user", "content": prompt}) st.session_state.chat_messages.append({"role": "user", "content": prompt}) placeholder = st.empty() random_seed = random.randint(0, 2 ** 32 - 1) setup_seed(random_seed) tools = [t for t in TOOLS if t['function']['name'] in st.session_state.get('selected_tools', [])] or None sys_prompt = [] if tools else [{"role": "system", "content": "你是MiniMind，一个乐于助人、知识渊博的AI助手。请用完整且友好的方式回答用户问题。"}] st.session_state.chat_messages = sys_prompt + st.session_state.chat_messages[-(st.session_state.history_chat_num + 1):] template_kwargs = {"tokenize": False, "add_generation_prompt": True} if st.session_state.get('enable_thinking', False): template_kwargs["open_thinking"] = True if tools: template_kwargs["tools"] = tools new_prompt = tokenizer.apply_chat_template(st.session_state.chat_messages, **template_kwargs) inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) generation_kwargs = { "input_ids": inputs.input_ids, "max_length": inputs.input_ids.shape[1] + st.session_state.max_new_tokens, "num_return_sequences": 1, "do_sample": True, "attention_mask": inputs.attention_mask, "pad_token_id": tokenizer.pad_token_id, "eos_token_id": tokenizer.eos_token_id, "temperature": st.session_state.temperature, "top_p": 0.85, "streamer": streamer, } Thread(target=model.generate, kwargs=generation_kwargs).start() answer = "" for new_text in streamer: answer += new_text placeholder.markdown(process_assistant_content(answer, is_streaming=True), unsafe_allow_html=True) full_answer = answer for _ in range(16): tool_calls = re.findall(r'(.*?)', answer, re.DOTALL) if not tool_calls: break st.session_state.chat_messages.append({"role": "assistant", "content": answer}) tool_results = [] for tc_str in tool_calls: try: tc = json.loads(tc_str.strip()) result = execute_tool(tc.get('name', ''), tc.get('arguments', {})) st.session_state.chat_messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)}) tool_results.append(f'
ToolCalled
{tc.get("name", "")}: {json.dumps(result, ensure_ascii=False)}
') except: pass full_answer += "\n" + "\n".join(tool_results) + "\n" placeholder.markdown(process_assistant_content(full_answer, is_streaming=True), unsafe_allow_html=True) new_prompt = tokenizer.apply_chat_template(st.session_state.chat_messages, **template_kwargs) inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) generation_kwargs["input_ids"] = inputs.input_ids generation_kwargs["attention_mask"] = inputs.attention_mask generation_kwargs["max_length"] = inputs.input_ids.shape[1] + st.session_state.max_new_tokens generation_kwargs["streamer"] = streamer Thread(target=model.generate, kwargs=generation_kwargs).start() answer = "" for new_text in streamer: answer += new_text placeholder.markdown(process_assistant_content(full_answer + answer, is_streaming=True), unsafe_allow_html=True) full_answer += answer answer = full_answer messages.append({"role": "assistant", "content": answer}) st.session_state.chat_messages.append({"role": "assistant", "content": answer}) if __name__ == "__main__": main()