Update inference_node.py

inference_node.py

@@ -13,91 +13,105 @@ from transformers import (
 # 1. 基础配置
 logging.basicConfig(level=logging.INFO, format="%(asctime)s-%(name)s-%(levelname)s-%(message)s")
 logger = logging.getLogger("inference_node_deepseek")
-app = FastAPI(title="推理节点服务（Deepseek-8B）")
+app = FastAPI(title="推理节点服务（DeepSeek-Coder-V2）")
 
-# 2. 模型配置（替换为Deepseek-8B-Instruct模型）
-# Deepseek-8B在数学推理和编程任务上表现突出，适合专业场景
-MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-llm-8b-instruct")  
-HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")  # 公开模型，无需令牌
+# 2. 模型配置：使用 Hugging Face 公开存在的 DeepSeek 模型
+# 正确 ID：deepseek-ai/deepseek-coder-v2（代码专用，公开无需令牌）
+MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-coder-v2")  
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")  # 公开模型，可留空
 
-# 3. 4bit量化配置（适配16G内存，优化Deepseek性能）
+# 3. 4bit量化配置（适配16G内存，DeepSeek 优化）
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16  # Deepseek更适合float16计算类型
+    bnb_4bit_compute_dtype=torch.float16  # 降低显存占用，适配 DeepSeek
 )
 
-# 4. 加载模型
+# 4. 加载 DeepSeek 模型（确保无 ID 错误）
 try:
     logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化）")
+    # 加载 Tokenizer（DeepSeek-Coder 专用配置）
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         token=HF_TOKEN,
         padding_side="right",
-        trust_remote_code=True  # Deepseek需要启用远程代码信任
+        trust_remote_code=True  # 必需：DeepSeek 模型需加载自定义代码
     )
-    # 设置pad_token（Deepseek默认没有pad_token）
+    # 手动设置 pad_token（DeepSeek 默认无，避免生成警告）
     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/CPU
         token=HF_TOKEN,
         trust_remote_code=True,
         torch_dtype=torch.float16
     )
-    # 流式生成器配置（保留必要的特殊标记）
+    # 流式生成器（保留代码格式所需的特殊标记）
     streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
     logger.info(f"模型 {MODEL_NAME} 加载成功！显存占用约 5-6GB（4bit 量化）")
 except Exception as e:
     logger.error(f"模型加载失败：{str(e)}", exc_info=True)
     raise SystemExit(f"服务终止：{str(e)}")
 
-# 5. 请求模型
+# 5. 请求模型（支持指定编程语言，适配代码生成场景）
 class NodeInferenceRequest(BaseModel):
-    prompt: str
+    prompt: str          # 代码需求（如“用Python写快速排序”）
+    language: str = "python"  # 可选：指定编程语言
     max_tokens: int = 1024
 
-# 6. 流式推理接口（适配Deepseek对话格式）
+# 6. 流式推理接口（适配 DeepSeek-Coder 对话格式）
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
     try:
-        # Deepseek要求的对话格式：[{"role": "user", "content": "..."}]
-        # 不需要额外添加系统提示，模型会自动处理
-        inputs = tokenizer.apply_chat_template(
-            [{"role": "user", "content": req.prompt}],
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt"
+        # 关键：DeepSeek-Coder 代码生成格式（明确语言类型，提升准确性）
+        code_prompt = f"""You are a professional code assistant. Write clean, runnable code for the following requirement.
+Programming Language: {req.language}
+Requirement: {req.prompt}
+Code (with comments):
+"""
+        # 构建输入（用标准 tokenize 方法，避免 build_chat_input 兼容问题）
+        inputs = tokenizer(
+            code_prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048  # 限制输入长度，预留生成空间
         ).to(model.device)
 
+        # 异步生成器
         async def generate_chunks():
             loop = asyncio.get_running_loop()
+            # 调用 DeepSeek-Coder 生成代码（低温度确保语法正确）
             outputs = await loop.run_in_executor(
                 None,
                 lambda: model.generate(
-                    inputs,
+                    **inputs,
                     streamer=streamer,
                     max_new_tokens=req.max_tokens,
                     do_sample=True,
-                    temperature=0.7,
+                    temperature=0.2,  # 代码生成用低温度（0.2-0.4），避免语法错误
                     top_p=0.95,
                     pad_token_id=tokenizer.pad_token_id,
                     eos_token_id=tokenizer.eos_token_id
                 )
             )
 
-            # 提取生成的部分（排除输入部分）
-            generated_tokens = outputs[0][len(inputs[0]):]
+            # 逐段解码代码（仅取生成部分，排除输入 Prompt）
+            generated_tokens = outputs[0][len(inputs["input_ids"][0]):]
             for token in generated_tokens:
                 if await request.is_disconnected():
+                    logger.info("客户端断开，停止生成")
                     break
+                # 解码 Token（保留代码格式）
                 token_text = tokenizer.decode(token, skip_special_tokens=True)
+                # 处理 JSON 转义（保留代码中的双引号和换行）
                 escaped_text = token_text.replace('"', '\\"').replace('\n', '\\n')
                 yield '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
+            # 生成结束标识
             yield '{"chunk":"","finish":true}\n'
 
         return StreamingResponse(generate_chunks(), media_type="application/x-ndjson")
@@ -107,17 +121,16 @@ async def stream_infer(req: NodeInferenceRequest, request: Request):
         logger.error(error_msg, exc_info=True)
         raise HTTPException(status_code=500, detail=error_msg)
 
-# 7. 健康检查
+# 7. 健康检查（确认模型正确加载）
 @app.get("/node/health")
 async def node_health():
     return {
         "status": "healthy",
         "model": MODEL_NAME,
         "support_stream": True,
-        "note": "Deepseek-8B 4bit量化，适配16G内存，擅长数学推理和编程任务"
+        "note": "DeepSeek-Coder-V2 4bit量化，适配16G内存，擅长Python/C++/Java代码生成"
     }
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
-    
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")