250426

scripts/web_demo.py

@@ -1,14 +1,13 @@
 import random
 import re
-import time
+from threading import Thread
 
+import torch
 import numpy as np
 import streamlit as st
-import torch
 
 st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed")
 
-# 在文件开头的 CSS 样式中修改按钮样式
 st.markdown("""
     <style>
         /* 添加操作按钮样式 */
@@ -70,7 +69,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def process_assistant_content(content):
-    if 'R1' not in MODEL_PATHS[selected_model][1]:
+    if model_source == "API" and 'R1' not in api_model_name:
+        return content
+    if model_source != "API" and 'R1' not in MODEL_PATHS[selected_model][1]:
         return content
 
     if '<think>' in content and '</think>' in content:
@@ -119,7 +120,6 @@ def init_chat_messages():
             if message["role"] == "assistant":
                 with st.chat_message("assistant", avatar=image_url):
                     st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
-                    # 在消息内容下方添加按钮
                     if st.button("🗑", key=f"delete_{i}"):
                         st.session_state.messages.pop(i)
                         st.session_state.messages.pop(i - 1)
@@ -137,8 +137,6 @@ def init_chat_messages():
 
     return st.session_state.messages
 
-
-# 添加这两个辅助函数
 def regenerate_answer(index):
     st.session_state.messages.pop()
     st.session_state.chat_messages.pop()
@@ -153,32 +151,34 @@ def delete_conversation(index):
     st.rerun()
 
 
-# 侧边栏模型选择
 st.sidebar.title("模型设定调整")
 
-st.sidebar.text("【注】训练数据偏差，增加上下文记忆时\n多轮对话（较单轮）容易出现能力衰减")
+# st.sidebar.text("训练数据偏差，增加上下文记忆时\n多轮对话（较单轮）容易出现能力衰减")
 st.session_state.history_chat_num = st.sidebar.slider("Number of Historical Dialogues", 0, 6, 0, step=2)
 # st.session_state.history_chat_num = 0
 st.session_state.max_new_tokens = st.sidebar.slider("Max Sequence Length", 256, 8192, 8192, step=1)
-st.session_state.top_p = st.sidebar.slider("Top-P", 0.8, 0.99, 0.85, step=0.01)
 st.session_state.temperature = st.sidebar.slider("Temperature", 0.6, 1.2, 0.85, step=0.01)
 
-# 模型路径映射
-MODEL_PATHS = {
-    "MiniMind2-R1 (0.1B)": ["../MiniMind2-R1", "MiniMind2-R1"],
-    "MiniMind2-Small-R1 (0.02B)": ["../MiniMind2-Small-R1", "MiniMind2-Small-R1"],
-    "MiniMind2 (0.1B)": ["../MiniMind2", "MiniMind2"],
-    "MiniMind2-MoE (0.15B)": ["../MiniMind2-MoE", "MiniMind2-MoE"],
-    "MiniMind2-Small (0.02B)": ["../MiniMind2-Small", "MiniMind2-Small"],
-    "MiniMind-V1 (0.1B)": ["../minimind-v1", "MiniMind-V1"],
-    "MiniMind-V1-MoE (0.1B)": ["../minimind-v1-moe", "MiniMind-V1-MoE"],
-    "MiniMind-V1-Small (0.02B)": ["../minimind-v1-small", "MiniMind-V1-Small"],
-}
+model_source = st.sidebar.radio("选择模型来源", ["本地模型", "API"], index=0)
 
-selected_model = st.sidebar.selectbox('Models', list(MODEL_PATHS.keys()), index=2)  # 默认选择 MiniMind2
-model_path = MODEL_PATHS[selected_model][0]
+if model_source == "API":
+    api_url = st.sidebar.text_input("API URL", value="http://127.0.0.1:8000/v1")
+    api_model_id = st.sidebar.text_input("Model ID", value="minimind")
+    api_model_name = st.sidebar.text_input("Model Name", value="MiniMind2")
+    api_key = st.sidebar.text_input("API Key", value="none", type="password")
+    slogan = f"Hi, I'm {api_model_name}"
+else:
+    MODEL_PATHS = {
+        "MiniMind2-R1 (0.1B)": ["../MiniMind2-R1", "MiniMind2-R1"],
+        "MiniMind2-Small-R1 (0.02B)": ["../MiniMind2-Small-R1", "MiniMind2-Small-R1"],
+        "MiniMind2 (0.1B)": ["../MiniMind2", "MiniMind2"],
+        "MiniMind2-MoE (0.15B)": ["../MiniMind2-MoE", "MiniMind2-MoE"],
+        "MiniMind2-Small (0.02B)": ["../MiniMind2-Small", "MiniMind2-Small"]
+    }
 
-slogan = f"Hi, I'm {MODEL_PATHS[selected_model][1]}"
+    selected_model = st.sidebar.selectbox('Models', list(MODEL_PATHS.keys()), index=2)  # 默认选择 MiniMind2
+    model_path = MODEL_PATHS[selected_model][0]
+    slogan = f"Hi, I'm {MODEL_PATHS[selected_model][1]}"
 
 image_url = "https://www.modelscope.cn/api/v1/studio/gongjy/MiniMind/repo?Revision=master&FilePath=images%2Flogo2.png&View=true"
 
@@ -205,23 +205,22 @@ def setup_seed(seed):
 
 
 def main():
-    model, tokenizer = load_model_tokenizer(model_path)
+    if model_source == "本地模型":
+        model, tokenizer = load_model_tokenizer(model_path)
+    else:
+        model, tokenizer = None, None
 
-    # 初始化消息列表
     if "messages" not in st.session_state:
         st.session_state.messages = []
         st.session_state.chat_messages = []
 
-    # Use session state messages
     messages = st.session_state.messages
 
-    # 在显示历史消息的循环中
     for i, message in enumerate(messages):
         if message["role"] == "assistant":
             with st.chat_message("assistant", avatar=image_url):
                 st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
                 if st.button("×", key=f"delete_{i}"):
-                    # 删除当前消息及其之后的所有消息
                     st.session_state.messages = st.session_state.messages[:i - 1]
                     st.session_state.chat_messages = st.session_state.chat_messages[:i - 1]
                     st.rerun()
@@ -230,14 +229,11 @@ def main():
                 f'<div style="display: flex; justify-content: flex-end;"><div style="display: inline-block; margin: 10px 0; padding: 8px 12px 8px 12px;  background-color: gray; border-radius: 10px; color:white; ">{message["content"]}</div></div>',
                 unsafe_allow_html=True)
 
-    # 处理新的输入或重新生成
     prompt = st.chat_input(key="input", placeholder="给 MiniMind 发送消息")
 
-    # 检查是否需要重新生成
     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  # 获取重新生成的位置
-        # 清除所有重新生成相关的状态
+        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')
@@ -246,48 +242,87 @@ def main():
         st.markdown(
             f'<div style="display: flex; justify-content: flex-end;"><div style="display: inline-block; margin: 10px 0; padding: 8px 12px 8px 12px;  background-color: gray; border-radius: 10px; color:white; ">{prompt}</div></div>',
             unsafe_allow_html=True)
-        messages.append({"role": "user", "content": prompt})
-        st.session_state.chat_messages.append({"role": "user", "content": prompt})
+        messages.append({"role": "user", "content": prompt[-st.session_state.max_new_tokens:]})
+        st.session_state.chat_messages.append({"role": "user", "content": prompt[-st.session_state.max_new_tokens:]})
 
         with st.chat_message("assistant", avatar=image_url):
             placeholder = st.empty()
-            random_seed = random.randint(0, 2 ** 32 - 1)
-            setup_seed(random_seed)
 
-            st.session_state.chat_messages = system_prompt + st.session_state.chat_messages[
-                                                             -(st.session_state.history_chat_num + 1):]
-            new_prompt = tokenizer.apply_chat_template(
-                st.session_state.chat_messages,
-                tokenize=False,
-                add_generation_prompt=True
-            )[-(st.session_state.max_new_tokens - 1):]
-
-            x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=device).unsqueeze(0)
-            with torch.no_grad():
-                res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=st.session_state.max_new_tokens,
-                                       temperature=st.session_state.temperature,
-                                       top_p=st.session_state.top_p, stream=True)
+            if model_source == "API":
                 try:
-                    for y in res_y:
-                        answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
-                        if (answer and answer[-1] == '�') or not answer:
-                            continue
+                    from openai import OpenAI
+
+                    client = OpenAI(
+                        api_key=api_key,
+                        base_url=api_url
+                    )
+                    history_num = st.session_state.history_chat_num + 1  # +1 是为了包含当前的用户消息
+                    conversation_history = system_prompt + st.session_state.chat_messages[-history_num:]
+                    answer = ""
+                    response = client.chat.completions.create(
+                        model=api_model_id,
+                        messages=conversation_history,
+                        stream=True,
+                        temperature=st.session_state.temperature
+                    )
+
+                    for chunk in response:
+                        content = chunk.choices[0].delta.content or ""
+                        answer += content
                         placeholder.markdown(process_assistant_content(answer), unsafe_allow_html=True)
-                except StopIteration:
-                    print("No answer")
 
-                assistant_answer = answer.replace(new_prompt, "")
-                messages.append({"role": "assistant", "content": assistant_answer})
-                st.session_state.chat_messages.append({"role": "assistant", "content": assistant_answer})
+                except Exception as e:
+                    answer = f"API调用出错: {str(e)}"
+                    placeholder.markdown(answer, unsafe_allow_html=True)
+            else:
+                random_seed = random.randint(0, 2 ** 32 - 1)
+                setup_seed(random_seed)
 
-                with st.empty():
-                    if st.button("×", key=f"delete_{len(messages) - 1}"):
-                        st.session_state.messages = st.session_state.messages[:-2]
-                        st.session_state.chat_messages = st.session_state.chat_messages[:-2]
-                        st.rerun()
+                st.session_state.chat_messages = system_prompt + st.session_state.chat_messages[
+                                                                 -(st.session_state.history_chat_num + 1):]
+                new_prompt = tokenizer.apply_chat_template(
+                    st.session_state.chat_messages,
+                    tokenize=False,
+                    add_generation_prompt=True
+                )
+
+                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), unsafe_allow_html=True)
+
+            messages.append({"role": "assistant", "content": answer})
+            st.session_state.chat_messages.append({"role": "assistant", "content": answer})
+            with st.empty():
+                if st.button("×", key=f"delete_{len(messages) - 1}"):
+                    st.session_state.messages = st.session_state.messages[:-2]
+                    st.session_state.chat_messages = st.session_state.chat_messages[:-2]
+                    st.rerun()
 
 
 if __name__ == "__main__":
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
     main()