Update app.py

app.py

@@ -1,51 +1,80 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware # 导入 CORS
+from pydantic import BaseModel, Field
+from typing import List, Optional
 
 # --- 1. 配置与模型加载 ---
 
+# 初始化 FastAPI 应用
+app = FastAPI(
+    title="Qwen 模型 API",
+    description="一个简单的API，用于与微调的Qwen模型进行交互，并可从任何网页调用。",
+    version="1.0.0"
+)
+
+# --- 新增：添加CORS中间件 ---
+# 这是允许浏览器JavaScript调用的关键改动。
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # 允许所有来源 (网站)
+    allow_credentials=True,
+    allow_methods=["*"],  # 允许所有方法 (GET, POST 等)
+    allow_headers=["*"],  # 允许所有请求头
+)
+
+# (文件的其余部分与之前相同)
+
 # 从环境变量或默认值加载模型ID
 MODEL_ID = os.getenv("MODEL_ID", "badanwang/teacher_basic_qwen3-0.6b")
 print(f"正在加载模型: {MODEL_ID}")
 
-# 加载分词器和模型
-# trust_remote_code=True 是加载Qwen等模型所必需的
-# device_map="auto" 会自动将模型分配到可用的硬件上（如GPU）
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype="auto",
-    device_map="auto",
-    trust_remote_code=True
-)
-print("模型加载成功！")
+# 使用一个全局字典来持有模型，避免重复加载
+model_objects = {}
+
+@app.on_event("startup")
+async def load_model():
+    """在应用启动时加载模型和分词器"""
+    print("应用启动... 开始加载模型...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True
+    )
+    model_objects['tokenizer'] = tokenizer
+    model_objects['model'] = model
+    print("模型加载成功！")
+
 
+# --- 2. 定义API的请求和响应数据结构 ---
 
-# --- 2. 核心推理函数 (API) ---
+class APIRequest(BaseModel):
+    prompt: str = Field(..., description="用户当前输入的问题。")
+    history: Optional[List[List[str]]] = Field(None, description="对话历史，格式为 [[user_msg_1, bot_msg_1], ...]。")
 
-def get_response(prompt: str, history: list[list[str]] = None):
-    """
-    一个简单的函数，用于与模型进行单次对话。
+class APIResponse(BaseModel):
+    response: str = Field(..., description="模型生成的回复。")
 
-    Args:
-        prompt (str): 用户当前输入的问题。
-        history (list[list[str]], optional): 对话历史，格式为 [[user_msg_1, bot_msg_1], ...]。默认为 None。
 
-    Returns:
-        str: 模型生成的回复。
-    """
+# --- 3. 核心推理函数 ---
+
+def get_response(prompt: str, history: Optional[List[List[str]]] = None) -> str:
+    tokenizer = model_objects['tokenizer']
+    model = model_objects['model']
+    
     if history is None:
         history = []
 
-    # 1. 构建消息列表
     messages = []
     for user_message, bot_message in history:
         messages.append({"role": "user", "content": user_message})
         messages.append({"role": "assistant", "content": bot_message})
     messages.append({"role": "user", "content": prompt})
 
-    # 2. 应用聊天模板并进行分词
-    # 这是与聊天模型正确交互的关键步骤
     input_ids = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
@@ -53,8 +82,6 @@ def get_response(prompt: str, history: list[list[str]] = None):
         return_tensors="pt"
     ).to(model.device)
 
-    # 3. 生成回复
-    # 这是一个阻塞式调用，会等待模型生成完毕
     outputs = model.generate(
         input_ids,
         max_new_tokens=1024,
@@ -63,32 +90,20 @@ def get_response(prompt: str, history: list[list[str]] = None):
         top_p=0.9
     )
     
-    # 4. 解码生成的文本
-    # `outputs[0]` 包含了输入的token和新生成的token，我们需要切片只获取新生成的部分
     response_ids = outputs[0][input_ids.shape[-1]:]
     response_text = tokenizer.decode(response_ids, skip_special_tokens=True)
 
     return response_text
 
-# --- 3. 使用示例 ---
-
-if __name__ == "__main__":
-    # 示例1: 单轮对话
-    print("\n--- 示例 1: 单轮对话 ---")
-    question1 = "你好，你是谁？"
-    print(f"用户: {question1}")
-    answer1 = get_response(question1)
-    print(f"模型: {answer1}")
-
-    # 示例2: 多轮对话
-    print("\n--- 示例 2: 多轮对话 ---")
-    # 首先，定义一个对话历史
-    chat_history = [
-        ["用Python写一个快速排序", "当然，这是快速排序的Python实现：\n```python\ndef quick_sort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quick_sort(left) + middle + quick_sort(right)\n\nprint(quick_sort())\n```"]
-    ]
-    question2 = "很好，你能解释一下它的工作原理吗？"
-    print(f"历史: {chat_history}")
-    print(f"用户: {question2}")
-    # 调用时传入历史记录
-    answer2 = get_response(question2, history=chat_history)
-    print(f"模型: {answer2}")
+
+# --- 4. 创建 API 端点 ---
+
+@app.post("/generate", response_model=APIResponse)
+async def generate(request: APIRequest):
+    """接收用户输入并返回模型的生成结果。"""
+    response_text = get_response(request.prompt, request.history)
+    return APIResponse(response=response_text)
+
+@app.get("/")
+def read_root():
+    return {"message": "欢迎使用Qwen模型API。请向 /generate 端点发送POST请求。"}