Update inference_node.py

inference_node.py

@@ -13,15 +13,13 @@ from transformers import (
 # 1. 基础配置
 logging.basicConfig(level=logging.INFO, format="%(asctime)s-%(name)s-%(levelname)s-%(message)s")
 logger = logging.getLogger("inference_node")
-app = FastAPI(title="推理节点服务（单一模型）")
+app = FastAPI(title="推理节点服务（Qwen-7B）")
 
-# 2. 模型配置（修复：使用正确的模型名，支持通过环境变量覆盖）
-# 正确模型名：Qwen/Qwen-0.5B-Instruct（Hugging Face 官方存在）
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen-0.5B-Instruct")  
-# 从环境变量获取 Hugging Face 令牌（必填，部分模型需登录）
-HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
+# 2. 模型配置（使用真实存在的 Qwen-7B）
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen-7B")  
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")  # Qwen-7B 公开，可留空
 
-# 3. 4bit量化配置（适配16G内存，降低显存占用）
+# 3. 4bit量化配置（适配16G内存）
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
@@ -29,104 +27,80 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# 4. 加载模型（修复：用 token 参数替代 use_auth_token，增加错误捕获）
+# 4. 加载模型
 try:
     logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化）")
-    # 加载 Tokenizer（修复参数：用 token 替代 use_auth_token）
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        token=HF_TOKEN,  # 新参数：传递 Hugging Face 令牌
-        padding_side="right",  # 避免生成时的警告
-        trust_remote_code=True  # 加载 Qwen 模型需开启（支持自定义代码）
+        token=HF_TOKEN,  # 公开模型可留空
+        padding_side="right",
+        trust_remote_code=True  # Qwen 模型必需
     )
-    # 加载量化模型
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         quantization_config=bnb_config,
-        device_map="auto",  # 自动分配 GPU/CPU（优先用 GPU）
-        token=HF_TOKEN,  # 传递令牌（部分模型需授权）
-        trust_remote_code=True  # Qwen 模型必需
+        device_map="auto",
+        token=HF_TOKEN,
+        trust_remote_code=True
     )
-    # 流式生成器（逐段输出结果）
     streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    logger.info(f"模型 {MODEL_NAME} 加载成功！显存占用约 2-3GB")
+    logger.info(f"模型 {MODEL_NAME} 加载成功！显存占用约 4-5GB（4bit 量化）")
 except Exception as e:
     logger.error(f"模型加载失败：{str(e)}", exc_info=True)
-    # 启动时加载失败直接退出（避免服务异常运行）
-    raise SystemExit(f"模型加载失败，服务终止：{str(e)}")
+    raise SystemExit(f"服务终止：{str(e)}")
 
-# 5. 请求模型（与总控约定的格式，无修改）
+# 5. 请求模型
 class NodeInferenceRequest(BaseModel):
-    prompt: str  # 总控拼接好的完整 Prompt（含用户上下文）
-    max_tokens: int = 1024  # 最大生成长度
+    prompt: str
+    max_tokens: int = 1024
 
-# 6. 流式推理接口（核心逻辑无修改，确保异步兼容）
+# 6. 流式推理接口
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
     try:
-        # 预处理 Prompt（Qwen 模型需用专用方法构建输入）
         inputs = tokenizer.build_chat_input(
-            [{"role": "user", "content": req.prompt}],  # 适配 Qwen 对话格式
-            add_generation_prompt=True  # 自动添加“助手回复”的提示
+            [{"role": "user", "content": req.prompt}],
+            add_generation_prompt=True
         ).to(model.device)
 
-        # 异步生成器：避免阻塞 FastAPI 事件循环
         async def generate_chunks():
-            generated_text = ""
-            # 用线程池执行同步的模型生成（避免阻塞异步接口）
             loop = asyncio.get_running_loop()
             outputs = await loop.run_in_executor(
-                None,  # 使用默认线程池
+                None,
                 lambda: model.generate(
                     **inputs,
                     streamer=streamer,
                     max_new_tokens=req.max_tokens,
                     do_sample=True,
-                    temperature=0.7,  # 随机性（0~1，越小越确定）
-                    pad_token_id=tokenizer.eos_token_id  # 避免警告
+                    temperature=0.7,
+                    pad_token_id=tokenizer.eos_token_id
                 )
             )
 
-            # 逐段解码并返回结果
             for token in outputs[0][len(inputs["input_ids"][0]):]:
-                # 检查客户端是否断开连接（提前终止，节省资源）
                 if await request.is_disconnected():
-                    logger.info("客户端断开连接，停止生成")
                     break
-                # 解码 Token（跳过特殊字符）
                 token_text = tokenizer.decode(token, skip_special_tokens=True)
-                generated_text += token_text
-                # 处理双引号转义（确保 JSON 格式合法）
                 escaped_text = token_text.replace('"', '\\"')
-                # 用 format 拼接 JSON，避免引号冲突
                 yield '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
-
-            # 生成结束标识
             yield '{"chunk":"","finish":true}\n'
 
-        # 返回流式响应（指定媒体类型为 JSON 流）
-        return StreamingResponse(
-            generate_chunks(),
-            media_type="application/x-ndjson",
-            headers={"Cache-Control": "no-cache"}
-        )
+        return StreamingResponse(generate_chunks(), media_type="application/x-ndjson")
 
     except Exception as e:
-        error_msg = f"推理失败：{str(e)}"
-        logger.error(error_msg, exc_info=True)
-        raise HTTPException(status_code=500, detail=error_msg)
+        logger.error(f"推理失败：{str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# 7. 健康检查接口（总控用于检测节点状态）
+# 7. 健康检查
 @app.get("/node/health")
 async def node_health():
     return {
         "status": "healthy",
         "model": MODEL_NAME,
         "support_stream": True,
-        "note": "Qwen-0.5B-Instruct 4bit量化，显存占用~2GB"
+        "note": "Qwen-7B 4bit量化，适配16G内存"
     }
 
 if __name__ == "__main__":
     import uvicorn
-    # 启动服务（Hugging Face Space 默认端口 7860）
     uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")