Update inference_node.py

inference_node.py

@@ -7,7 +7,7 @@ import torch
 import asyncio
 from transformers import (
     AutoModelForCausalLM, AutoTokenizer,
-    BitsAndBytesConfig
+    BitsAndBytesConfig, TextStreamer
 )
 
 # 1. 基础配置
@@ -27,7 +27,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# 4. 加载模型（移除TextStreamer，避免兼容性问题）
+# 4. 加载模型（关键：显式处理tokenizer缺失的配置）
 try:
     logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化）")
     tokenizer = AutoTokenizer.from_pretrained(
@@ -35,7 +35,7 @@ try:
         token=HF_TOKEN,
         padding_side="right",  # 右侧padding，避免生成时截断
         trust_remote_code=True,  # Qwen模型必需（加载自定义tokenizer）
-        eos_token="<|endoftext|>",  # 显式指定结束符
+        eos_token="<|endoftext|>",  # 显式指定结束符（兼容旧版本）
         pad_token="<|endoftext|>"   # 显式指定padding符（避免生成警告）
     )
     model = AutoModelForCausalLM.from_pretrained(
@@ -46,8 +46,8 @@ try:
         trust_remote_code=True,
         torch_dtype=torch.bfloat16  # 匹配量化计算精度
     )
-    # 关键：设置模型为评估模式（避免训练模式的随机行为）
-    model.eval()
+    # 流式输出配置（跳过提示词，只返回生成内容）
+    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)
@@ -58,122 +58,66 @@ class NodeInferenceRequest(BaseModel):
     prompt: str  # 用户提问内容
     max_tokens: int = 1024  # 最大生成长度（默认1024）
     temperature: float = 0.7  # 随机性（0-1，越大越多样）
-    top_p: float = 0.9  # 采样Top-p（配合temperature使用）
 
-# 6. 流式推理接口（核心：手动逐token生成，显式管理past_key_values）
+# 6. 流式推理接口（核心修复：绕开chat_template，直接构建输入）
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
     try:
         # --------------------------
-        # 1. 构建Qwen原生对话输入格式
+        # 关键修复：手动构建Qwen原生对话格式
+        # Qwen要求格式：<|user|>用户输入<|end|><|assistant|>
         # --------------------------
         user_prompt = req.prompt.strip()
-        if not user_prompt:
-            raise ValueError("用户输入prompt不能为空")
-        
-        # Qwen-7B 原生对话格式：<|user|>提问<|end|><|assistant|>
+        # 构建模型能理解的输入文本（无需依赖chat_template）
         input_text = f"<|user|>{user_prompt}<|end|><|assistant|>"
         
-        # 编码输入（转换为GPU张量，不进行padding，避免冗余）
+        # 编码输入（转换为模型可处理的张量，并移动到GPU）
         inputs = tokenizer(
             input_text,
             return_tensors="pt",  # 返回PyTorch张量
             truncation=True,      # 截断过长输入（避免OOM）
-            max_length=2048       # 输入最大长度（匹配模型能力）
+            max_length=2048       # 输入最大长度（根据模型能力调整）
         ).to(model.device)
 
-        # --------------------------
-        # 2. 初始化生成状态（显式管理past_key_values）
-        # --------------------------
-        input_ids = inputs["input_ids"]  # 初始输入token
-        attention_mask = inputs["attention_mask"]  # 初始注意力掩码
-        past_key_values = None  # 初始化键值对缓存（后续逐token更新）
-        generated_tokens = []  # 记录已生成的token（用于最终校验）
-        max_new_tokens = min(req.max_tokens, 2048)  # 限制最大生成长度（避免OOM）
-
-        # --------------------------
-        # 3. 异步流式生成（手动逐token生成，避免TextStreamer问题）
-        # --------------------------
+        # 异步生成流式内容（避免阻塞FastAPI主线程）
         async def generate_chunks():
-            nonlocal past_key_values, generated_tokens
             loop = asyncio.get_running_loop()
-
-            for _ in range(max_new_tokens):
+            # 在线程池中运行同步的模型生成（不阻塞事件循环）
+            outputs = await loop.run_in_executor(
+                None,  # 使用默认线程池
+                lambda: model.generate(
+                    **inputs,
+                    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（生成停止标志）
+                )
+            )
+
+            # 提取生成的内容（排除输入部分，只取新生成的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("客户端已断开连接，停止生成")
-                    yield '{"chunk":"","finish":true}\n'
-                    return
-
-                # --------------------------
-                # 关键：在线程池中运行模型生成（每次只生成1个token）
-                # --------------------------
-                try:
-                    # 显式传入past_key_values，确保缓存不丢失
-                    generate_kwargs = {
-                        "input_ids": input_ids,
-                        "attention_mask": attention_mask,
-                        "past_key_values": past_key_values,
-                        "max_new_tokens": 1,  # 每次只生成1个token（流式核心）
-                        "do_sample": True,
-                        "temperature": req.temperature,
-                        "top_p": req.top_p,
-                        "pad_token_id": tokenizer.pad_token_id,
-                        "eos_token_id": tokenizer.eos_token_id,
-                        "use_cache": True,  # 必须开启缓存，否则past_key_values无效
-                        "return_dict_in_generate": True,  # 返回字典格式，便于获取past_key_values
-                        "output_scores": False  # 关闭分数输出，减少计算开销
-                    }
-
-                    # 同步生成（在独立线程中运行，不阻塞FastAPI事件循环）
-                    outputs = await loop.run_in_executor(
-                        None,  # 使用默认线程池
-                        lambda: model.generate(**generate_kwargs)
-                    )
-
-                    # --------------------------
-                    # 4. 更新生成状态（关键：保存past_key_values）
-                    # --------------------------
-                    next_token = outputs.sequences[:, -1:]  # 获取最新生成的1个token
-                    past_key_values = outputs.past_key_values  # 更新缓存（避免下一轮为None）
-                    generated_tokens.append(next_token.item())  # 记录生成的token
-
-                    # --------------------------
-                    # 5. 解码token并返回（流式输出）
-                    # --------------------------
-                    # 解码单个token（跳过特殊符号，清理空格）
-                    token_text = tokenizer.decode(
-                        next_token[0],  # 取batch中的第一个（仅单条请求）
-                        skip_special_tokens=True,
-                        clean_up_tokenization_spaces=True
-                    )
-
-                    # 转义双引号（避免JSON格式错误）
-                    escaped_text = token_text.replace('"', '\\"')
-                    # 按NDJSON格式返回（每行一个JSON对象，客户端可逐行解析）
-                    yield f'{{"chunk":"{escaped_text}","finish":false}}\n'
-
-                    # 检查是否生成结束符（eos_token），是的话终止生成
-                    if next_token.item() == tokenizer.eos_token_id:
-                        logger.info(f"生成结束符（eos_token），停止生成")
-                        break
-
-                    # 更新下一轮的输入（仅使用最新生成的token，减少计算量）
-                    input_ids = next_token
-                    # 扩展注意力掩码（新token的掩码为1）
-                    attention_mask = torch.cat([
-                        attention_mask,
-                        torch.ones((1, 1), device=model.device, dtype=torch.long)
-                    ], dim=-1)
-
-                except Exception as gen_e:
-                    logger.error(f"逐token生成失败：{str(gen_e)}", exc_info=True)
-                    yield f'{{"chunk":"生成过程异常：{str(gen_e)}","finish":true}}\n'
-                    return
-
-            # --------------------------
-            # 6. 生成结束（返回终止标志）
-            # --------------------------
+                    break
+                # 解码单个token（跳过特殊符号，如<|end|>）
+                token_text = tokenizer.decode(
+                    token,
+                    skip_special_tokens=True,  # 跳过特殊token（如结束符、分隔符）
+                    clean_up_tokenization_spaces=True  # 清理多余空格
+                )
+                # 转义双引号（避免JSON格式错误）
+                escaped_text = token_text.replace('"', '\\"')
+                # 按NDJSON格式返回（每行一个JSON对象，兼容流式解析）
+                yield f'{{"chunk":"{escaped_text}","finish":false}}\n'
+            
+            # 生成结束标志（告知客户端生成完成）
             yield '{"chunk":"","finish":true}\n'
 
         # 返回流式响应（媒体类型为application/x-ndjson，支持逐行解析）
@@ -183,9 +127,8 @@ async def stream_infer(req: NodeInferenceRequest, request: Request):
         )
 
     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)}", exc_info=True)  # 记录详细错误堆栈
+        raise HTTPException(status_code=500, detail=f"推理服务异常：{str(e)}")
 
 # 7. 健康检查接口（用于监控服务状态）
 @app.get("/node/health")
@@ -196,18 +139,18 @@ async def node_health():
         "status": "healthy" if is_model_ready else "unhealthy",
         "model": MODEL_NAME,
         "support_stream": True,
-        "note": "Qwen-7B 4bit量化（手动逐token生成，解决past_key_values问题）",
+        "note": "Qwen-7B 4bit量化（适配16G内存），绕开chat_template兼容旧版本",
         "timestamp": str(asyncio.get_event_loop().time())
     }
 
 # 8. 启动服务（仅在直接运行脚本时执行）
 if __name__ == "__main__":
     import uvicorn
-    # 启动UVicorn服务（单进程，避免模型重复加载）
+    # 启动UVicorn服务（host=0.0.0.0允许外部访问，port=7860为默认端口）
     uvicorn.run(
         app,
         host="0.0.0.0",
         port=7860,
         log_level="info",
-        workers=1  # 模型不支持多进程共享，必须设为1
+        workers=1  # 单进程（模型不支持多进程共享，避免重复加载）
     )