Update app.py

app.py

@@ -1,22 +1,56 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import warnings
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import random
+import numpy as np
 
-# 加载模型和tokenizer
-model_name = "cmz1024/minimind-zero"  # 替换为你的模型路径
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-model.eval()
+warnings.filterwarnings('ignore')
 
-if torch.cuda.is_available():
-    model = model.cuda()
+# 设置可复现的随机种子
+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 generate_text(prompt, max_length=512, temperature=0.7, top_p=0.9):
-    # 对输入进行编码
-    inputs = tokenizer(prompt, return_tensors="pt")
+# 加载模型和分词器
+model_path = "cmz1024/minimind-zero"  # 替换为你的模型路径
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
+
+# 将模型移至GPU（如果可用）
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device).eval()
+
+print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M(illion)')
+
+# 生成文本函数
+def generate_text(prompt, max_length=512, temperature=0.85, top_p=0.85, history_cnt=0):
+    # 设置随机种子
+    setup_seed(random.randint(0, 2048))
+    
+    # 处理历史对话
+    messages = []
+    if history_cnt > 0 and 'chat_history' in globals():
+        messages = chat_history[-history_cnt:] if len(chat_history) > 0 else []
+    
+    # 添加当前用户输入
+    messages.append({"role": "user", "content": prompt})
+    
+    # 应用聊天模板
+    new_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
     
-    if torch.cuda.is_available():
-        inputs = {k: v.cuda() for k, v in inputs.items()}
+    # 对输入进行编码
+    inputs = tokenizer(new_prompt, return_tensors="pt").to(device)
+    input_length = inputs["input_ids"].shape[1]
     
     # 生成文本
     with torch.no_grad():
@@ -26,26 +60,62 @@ def generate_text(prompt, max_length=512, temperature=0.7, top_p=0.9):
             temperature=temperature,
             top_p=top_p,
             pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
         )
     
-    # 解码输出
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # 只解码新生成的部分
+    generated_text = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+    
+    # 更新对话历史
+    if 'chat_history' in globals():
+        chat_history.append({"role": "user", "content": prompt})
+        chat_history.append({"role": "assistant", "content": generated_text})
+    
     return generated_text
 
+# 初始化全局对话历史
+chat_history = []
+
+# 清除对话历史的函数
+def clear_history():
+    global chat_history
+    chat_history = []
+    return "对话历史已清除"
+
 # 创建Gradio界面
-demo = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(label="输入提示词", lines=3),
-        gr.Slider(minimum=1, maximum=1024, value=512, label="最大生成长度"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, label="Top-p"),
-    ],
-    outputs=gr.Textbox(label="生成结果", lines=10),
-    title="MiniMind 文本生成",
-    description="一个简单的文本生成demo"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# MiniMind 模型演示")
+    
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(label="输入", placeholder="请输入您的问题...", lines=5)
+            
+            with gr.Row():
+                submit_btn = gr.Button("提交")
+                clear_btn = gr.Button("清除历史")
+            
+            with gr.Accordion("高级选项", open=False):
+                max_length = gr.Slider(minimum=10, maximum=2048, value=512, step=1, label="最大生成长度")
+                temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.85, step=0.01, label="温度")
+                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.85, step=0.01, label="Top-p")
+                history_cnt = gr.Slider(minimum=0, maximum=10, value=0, step=2, label="历史对话轮数")
+        
+        with gr.Column():
+            output_text = gr.Textbox(label="输出", lines=25)
+            clear_output = gr.Textbox(label="状态", visible=True)
+    
+    # 设置事件
+    submit_btn.click(
+        fn=generate_text,
+        inputs=[input_text, max_length, temperature, top_p, history_cnt],
+        outputs=output_text
+    )
+    
+    clear_btn.click(
+        fn=clear_history,
+        inputs=[],
+        outputs=clear_output
+    )
+
+# 启动应用
+demo.launch()