Update app.py

app.py

@@ -1,41 +1,86 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# 修改为你上传的模型文件夹路径
-model_path = "tosei0000/code-AI"
+# 模型路径（本地或 huggingface repo）
+model_path = "tosei0000/code-AI"  # ← 修改为你的模型目录或 huggingface repo_id
 
-# 加载模型和分词器
+# 加载 tokenizer 和模型
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_path, 
-    trust_remote_code=True, 
-    torch_dtype=torch.float32  # 若使用 GPU，可改为 torch.float16
-)
+model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.float32)
 model.eval()
 
-def chat(prompt, max_new_tokens=512):
-    inputs = tokenizer(prompt, return_tensors="pt")
+# 系统人格设定，可修改为你喜欢的角色
+SYSTEM_PROMPT = (
+    "你是一个善良、聪明、幽默的 AI 编程助手，能写各种 Python/HTML/JavaScript 代码，"
+    "也能像朋友一样聊天。请保持耐心、有趣，尽可能详细地回答问题。\n"
+)
+
+# 对话历史记录（最多保存 5 轮）
+chat_history = []
+
+# 生成回复函数
+def generate_reply(user_input, chat_mode, max_tokens=512, temperature=0.7):
+    global chat_history
+
+    # 更新历史上下文
+    chat_history.append(f"用户: {user_input}")
+    if len(chat_history) > 5:
+        chat_history = chat_history[-5:]  # 只保留最近 5 条
+
+    # 拼接 prompt
+    if chat_mode == "代码生成":
+        prompt = f"{SYSTEM_PROMPT}\n请根据以下需求生成代码：\n{user_input}\n"
+    else:  # 聊天模式
+        prompt = SYSTEM_PROMPT + "\n" + "\n".join(chat_history) + "\n助手:"
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
             top_p=0.95,
-            temperature=0.8,
-            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id,
         )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response[len(prompt):]  # 返回去除原始prompt后的回答部分
-
-# 创建 Gradio 界面
-gr.Interface(
-    fn=chat,
-    inputs=[
-        gr.Textbox(label="你的提问", lines=3, placeholder="请输入 prompt..."),
-        gr.Slider(128, 1024, step=64, value=512, label="最大生成长度")
-    ],
-    outputs="text",
-    title="🧠 DeepSeek-R1 Chat Demo",
-    description="使用你本地上传的 DeepSeek-R1 模型运行的聊天机器人。"
-).launch()
+
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # 提取最后一句助手的回复
+    reply = decoded_output.split("助手:")[-1].strip()
+
+    # 保存助手回复到历史中
+    chat_history.append(f"助手: {reply}")
+
+    return reply
+
+# 重置历史按钮
+def reset_memory():
+    global chat_history
+    chat_history = []
+    return "记忆已重置。"
+
+# Gradio 界面
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 智能代码助理 + 聊天机器人")
+    gr.Markdown("支持代码生成与聊天功能，可记忆上下文，具备人格设定！")
+
+    with gr.Row():
+        chat_mode = gr.Radio(["聊天", "代码生成"], value="代码生成", label="对话模式")
+        reset_btn = gr.Button("🧹 重置记忆")
+
+    user_input = gr.Textbox(label="你的输入", lines=6, placeholder="输入代码需求或聊天内容...")
+    max_tokens = gr.Slider(50, 1024, value=512, step=10, label="最大生成长度")
+    temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="生成多样性（temperature）")
+
+    output = gr.Textbox(label="AI 回复", lines=10)
+
+    submit_btn = gr.Button("✨ 生成")
+
+    submit_btn.click(fn=generate_reply, inputs=[user_input, chat_mode, max_tokens, temperature], outputs=output)
+    reset_btn.click(fn=reset_memory, outputs=output)
+
+# 启动服务
+if __name__ == "__main__":
+    demo.launch()