Update app.py

app.py

@@ -1,86 +1,23 @@
-import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# 模型路径（本地或 huggingface repo）
-model_path = "tosei0000/code-AI"  # ← 修改为你的模型目录或 huggingface repo_id
+# 设置模型路径（替换为你上传后显示的路径名）
+model_path = "tosei0000/code-AI"  # 修改为你上传的文件夹名
 
-# 加载 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)
-model.eval()
+model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
 
-# 系统人格设定，可修改为你喜欢的角色
-SYSTEM_PROMPT = (
-    "你是一个善良、聪明、幽默的 AI 编程助手，能写各种 Python/HTML/JavaScript 代码，"
-    "也能像朋友一样聊天。请保持耐心、有趣，尽可能详细地回答问题。\n"
-)
+# 如果你有 GPU（Kaggle 支持 GPU），把模型放到 GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-# 对话历史记录（最多保存 5 轮）
-chat_history = []
+# 推理函数
+def chat(prompt, max_new_tokens=100):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 生成回复函数
-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_tokens,
-            temperature=temperature,
-            top_p=0.95,
-            do_sample=True,
-            eos_token_id=tokenizer.eos_token_id,
-        )
-
-    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()
+# 测试
+response = chat("你好，请介绍一下你自己。")
+print(response)