Update inference_node.py

inference_node.py

@@ -15,11 +15,11 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s-%(name)s-%(levelname
 logger = logging.getLogger("inference_node")
 app = FastAPI(title="推理节点服务（Qwen-7B）")
 
-# 2. 模型配置
+# 2. 模型配置（Qwen-7B 公开模型，无需HF Token）
 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen-7B")  
-HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN", "")  # 留空即可
 
-# 3. 4bit量化配置
+# 3. 4bit量化配置（适配16G内存，显存占用约4-5GB）
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
@@ -27,102 +27,130 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# 4. 加载模型
+# 4. 加载模型（关键：显式处理tokenizer缺失的配置）
 try:
     logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化）")
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         token=HF_TOKEN,
-        padding_side="right",
-        trust_remote_code=True
+        padding_side="right",  # 右侧padding，避免生成时截断
+        trust_remote_code=True,  # Qwen模型必需（加载自定义tokenizer）
+        eos_token="<|endoftext|>",  # 显式指定结束符（兼容旧版本）
+        pad_token="<|endoftext|>"   # 显式指定padding符（避免生成警告）
     )
-    # 添加缺失的pad_token（如果需要）
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-        
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         quantization_config=bnb_config,
-        device_map="auto",
+        device_map="auto",  # 自动分配设备（优先GPU）
         token=HF_TOKEN,
-        trust_remote_code=True
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16  # 匹配量化计算精度
     )
+    # 流式输出配置（跳过提示词，只返回生成内容）
     streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     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)}")
 
-# 5. 请求模型
+# 5. 请求体定义（用户输入prompt和生成参数）
 class NodeInferenceRequest(BaseModel):
-    prompt: str
-    max_tokens: int = 1024
+    prompt: str  # 用户提问内容
+    max_tokens: int = 1024  # 最大生成长度（默认1024）
+    temperature: float = 0.7  # 随机性（0-1，越大越多样）
 
-# 6. 流式推理接口 - 修复了build_chat_input的问题
+# 6. 流式推理接口（核心修复：绕开chat_template，直接构建输入）
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
     try:
-        # 构建对话历史 - 替代build_chat_input方法
-        conversation = [
-            {"from": "user", "value": req.prompt},
-            {"from": "assistant", "value": ""}  # 空的assistant回复作为生成提示
-        ]
-        
-        # 将对话转换为模型输入
-        input_text = tokenizer.apply_chat_template(
-            conversation,
-            tokenize=False,
-            add_generation_prompt=True
-        )
+        # --------------------------
+        # 关键修复：手动构建Qwen原生对话格式
+        # Qwen要求格式：<|user|>用户输入<|end|><|assistant|>
+        # --------------------------
+        user_prompt = req.prompt.strip()
+        # 构建模型能理解的输入文本（无需依赖chat_template）
+        input_text = f"<|user|>{user_prompt}<|end|><|assistant|>"
         
-        # 编码输入
+        # 编码输入（转换为模型可处理的张量，并移动到GPU）
         inputs = tokenizer(
             input_text,
-            return_tensors="pt",
-            padding=True,
-            truncation=True
+            return_tensors="pt",  # 返回PyTorch张量
+            truncation=True,      # 截断过长输入（避免OOM）
+            max_length=2048       # 输入最大长度（根据模型能力调整）
         ).to(model.device)
 
+        # 异步生成流式内容（避免阻塞FastAPI主线程）
         async def generate_chunks():
             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,
-                    pad_token_id=tokenizer.eos_token_id
+                    streamer=streamer,          # 流式输出支持
+                    max_new_tokens=req.max_tokens,  # 最大生成长度
+                    do_sample=True,             # 启��采样（生成多样内容）
+                    temperature=req.temperature,  # 随机性控制
+                    pad_token_id=tokenizer.pad_token_id,  # padding符ID
+                    eos_token_id=tokenizer.eos_token_id   # 结束符ID（生成停止标志）
                 )
             )
 
-            # 提取生成的部分（排除输入部分）
-            generated_tokens = outputs[0][len(inputs["input_ids"][0]):]
+            # 提取生成的内容（排除输入部分，只取新生成的token）
+            input_token_len = inputs["input_ids"].shape[1]  # 输入token长度
+            generated_tokens = outputs[0][input_token_len:]  # 仅保留新生成的token
+
+            # 逐token解码并返回（流式输出核心）
             for token in generated_tokens:
+                # 检查客户端是否断开连接（避免无效生成）
                 if await request.is_disconnected():
+                    logger.info("客户端已断开连接，停止生成")
                     break
-                token_text = tokenizer.decode(token, skip_special_tokens=True)
+                # 解码单个token（跳过特殊符号，如<|end|>）
+                token_text = tokenizer.decode(
+                    token,
+                    skip_special_tokens=True,  # 跳过特殊token（如结束符、分隔符）
+                    clean_up_tokenization_spaces=True  # 清理多余空格
+                )
+                # 转义双引号（避免JSON格式错误）
                 escaped_text = token_text.replace('"', '\\"')
-                yield '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
+                # 按NDJSON格式返回（每行一个JSON对象，兼容流式解析）
+                yield f'{{"chunk":"{escaped_text}","finish":false}}\n'
+            
+            # 生成结束标志（告知客户端生成完成）
             yield '{"chunk":"","finish":true}\n'
 
-        return StreamingResponse(generate_chunks(), media_type="application/x-ndjson")
+        # 返回流式响应（媒体类型为application/x-ndjson，支持逐行解析）
+        return StreamingResponse(
+            generate_chunks(),
+            media_type="application/x-ndjson"
+        )
 
     except Exception as e:
-        logger.error(f"推理失败：{str(e)}", exc_info=True)  # 添加详细异常信息
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"推理失败：{str(e)}", exc_info=True)  # 记录详细错误堆栈
+        raise HTTPException(status_code=500, detail=f"推理服务异常：{str(e)}")
 
-# 7. 健康检查
+# 7. 健康检查接口（用于监控服务状态）
 @app.get("/node/health")
 async def node_health():
+    # 检查模型和tokenizer是否正常加载
+    is_model_ready = model is not None and tokenizer is not None
     return {
-        "status": "healthy",
+        "status": "healthy" if is_model_ready else "unhealthy",
         "model": MODEL_NAME,
         "support_stream": True,
-        "note": "Qwen-7B 4bit量化，适配16G内存"
+        "note": "Qwen-7B 4bit量化（适配16G内存），绕开chat_template兼容旧版本",
+        "timestamp": str(asyncio.get_event_loop().time())
     }
 
+# 8. 启动服务（仅在直接运行脚本时执行）
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
+    # 启动UVicorn服务（host=0.0.0.0允许外部访问，port=7860为默认端口）
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,
+        log_level="info",
+        workers=1  # 单进程（模型不支持多进程共享，避免重复加载）
+    )