Update inference_node.py

inference_node.py

@@ -13,11 +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_deepseek")
-app = FastAPI(title="推理节点服务（DeepSeek-Coder-V2）")
+app = FastAPI(title="推理节点服务（DeepSeek-Math-7B-RL）")
 
-# 2. 模型配置：使用 Hugging Face 公开存在的 DeepSeek 模型
-# 正确 ID：deepseek-ai/deepseek-coder-v2（代码专用，公开无需令牌）
-MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-coder-v2")  
+# 2. 模型配置：使用 DeepSeek 官方公开且无访问限制的模型
+# 正确 ID：deepseek-ai/deepseek-math-7b-rl（公开无需令牌，支持数学/通用对话）
+# 新增 revision="main"：明确加载主分支，避免版本解析错误
+MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-math-7b-rl")  
+MODEL_REVISION = "main"  # 关键：指定模型分支，确保找到文件
 HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")  # 公开模型，可留空
 
 # 3. 4bit量化配置（适配16G内存，DeepSeek 优化）
@@ -28,64 +30,70 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16  # 降低显存占用，适配 DeepSeek
 )
 
-# 4. 加载 DeepSeek 模型（确保无 ID 错误）
+# 4. 加载 DeepSeek 模型（新增 revision 参数，确保找到文件）
 try:
-    logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化）")
-    # 加载 Tokenizer（DeepSeek-Coder 专用配置）
+    logger.info(f"开始加载模型：{MODEL_NAME}（分支：{MODEL_REVISION}，4bit量化）")
+    # 加载 Tokenizer（新增 revision 参数，匹配模型文件）
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
+        revision=MODEL_REVISION,  # 关键：指定分支
         token=HF_TOKEN,
         padding_side="right",
-        trust_remote_code=True  # 必需：DeepSeek 模型需加载自定义代码
+        trust_remote_code=True  # DeepSeek 必需：加载自定义 Tokenizer 逻辑
     )
     # 手动设置 pad_token（DeepSeek 默认无，避免生成警告）
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     
-    # 加载量化模型
+    # 加载量化模型（同样指定 revision）
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
+        revision=MODEL_REVISION,  # 关键：与 Tokenizer 分支一致
         quantization_config=bnb_config,
         device_map="auto",  # 自动分配 GPU/CPU
         token=HF_TOKEN,
-        trust_remote_code=True,
+        trust_remote_code=True,  # DeepSeek 必需：加载自定义模型结构
         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          # 代码需求（如“用Python写快速排序”）
-    language: str = "python"  # 可选：指定编程语言
+    prompt: str          # 输入需求（如“解一元二次方程 x²-5x+6=0”）
     max_tokens: int = 1024
+    is_math: bool = False  # 可选：是否为数学任务，优化生成逻辑
 
-# 6. 流式推理接口（适配 DeepSeek-Coder 对话格式）
+# 6. 流式推理接口（适配 DeepSeek 对话格式，支持数学场景）
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
     try:
-        # 关键：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 兼容问题）
+        # 适配 DeepSeek 对话格式（数学任务添加特殊提示，提升准确性）
+        if req.is_math:
+            prompt = f"""你是专业的数学助手，需详细步骤解答数学问题。
+问题：{req.prompt}
+解答（含步骤）："""
+        else:
+            prompt = f"""你是通用对话助手，需清晰、准确地回答问题。
+问题：{req.prompt}
+回答："""
+        
+        # 构建输入（用标准 tokenize 方法，避免兼容问题）
         inputs = tokenizer(
-            code_prompt,
+            prompt,
             return_tensors="pt",
             truncation=True,
             max_length=2048  # 限制输入长度，预留生成空间
         ).to(model.device)
 
-        # 异步生成器
+        # 异步生成器（确保流式输出）
         async def generate_chunks():
             loop = asyncio.get_running_loop()
-            # 调用 DeepSeek-Coder 生成代码（低温度确保语法正确）
+            # 调用 DeepSeek 生成（数学任务用低温度，确保步骤正确）
             outputs = await loop.run_in_executor(
                 None,
                 lambda: model.generate(
@@ -93,22 +101,24 @@ Code (with comments):
                     streamer=streamer,
                     max_new_tokens=req.max_tokens,
                     do_sample=True,
-                    temperature=0.2,  # 代码生成用低温度（0.2-0.4），避免语法错误
+                    temperature=0.3 if req.is_math else 0.7,  # 数学任务低温度（0.3）
                     top_p=0.95,
                     pad_token_id=tokenizer.pad_token_id,
                     eos_token_id=tokenizer.eos_token_id
                 )
             )
 
-            # 逐段解码代码（仅取生成部分，排除输入 Prompt）
+            # 逐段解码（仅取生成部分，排除输入 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（跳过结束符，保留纯文本）
                 token_text = tokenizer.decode(token, skip_special_tokens=True)
-                # 处理 JSON 转义（保留代码中的双引号和换行）
+                if token_text.endswith(tokenizer.eos_token):
+                    break
+                # 处理 JSON 转义（确保总控能解析）
                 escaped_text = token_text.replace('"', '\\"').replace('\n', '\\n')
                 yield '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
             # 生成结束标识
@@ -127,8 +137,9 @@ async def node_health():
     return {
         "status": "healthy",
         "model": MODEL_NAME,
+        "model_revision": MODEL_REVISION,
         "support_stream": True,
-        "note": "DeepSeek-Coder-V2 4bit量化，适配16G内存，擅长Python/C++/Java代码生成"
+        "note": "DeepSeek-Math-7B-RL 4bit量化，适配16G内存，支持数学推理和通用对话"
     }
 
 if __name__ == "__main__":