本地大模型落地医疗场景：Qwen3-32B与ChatGLM3-6B高血压慢病筛查实战测评19.9

import os
# 禁用vLLM v1引擎（v1对bitsandbytes支持不完善，回退v0稳定引擎）
os.environ["VLLM_USE_V1"] = "0"
# 屏蔽vLLM INFO/WARNING日志，只显示ERROR及以上
os.environ["VLLM_LOGGING_LEVEL"] = "ERROR"

from vllm import LLM, SamplingParams
from modelscope import snapshot_download
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import List, Literal
import subprocess
import time
import uuid
import re
import json
import uvicorn

# ===================== 1. 配置参数 =====================
model_name = "qwen/Qwen3-32B"
cache_dir = "/home/model"

# ===================== 2. 下载模型到本地 =====================
print("正在下载/校验模型缓存...")
local_model_path = snapshot_download(
    model_name,
    cache_dir=cache_dir,
    revision="master"
)
print(f"模型已就绪，本地路径：{local_model_path}")

# ===================== 3. vLLM加载模型（4bit量化） =====================
print("正在加载模型（vLLM + bitsandbytes 4bit量化）...")
llm = LLM(
    model=local_model_path,
    trust_remote_code=True,
    # 4bit量化（运行时量化FP16权重）
    quantization="bitsandbytes",
    load_format="bitsandbytes",
    dtype="bfloat16",
    # 显存利用率（拉满，32B-4bit单卡显存紧张）
    gpu_memory_utilization=0.97,
    # 最大上下文长度（筛查场景需容纳11项JSON输出）
    max_model_len=4096,
    # 无flash-attn时使用eager模式（如已装flash-attn可去掉此行）
    enforce_eager=True,
)

# 获取tokenizer（后续Chat接口复用）
tokenizer = llm.get_tokenizer()

# 输出模型加载后显存占用（vLLM在子进程中分配显存，需通过nvidia-smi读取）
import subprocess
def get_gpu_memory():
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits"],
        capture_output=True, text=True
    )
    used, total = result.stdout.strip().split(", ")
    return float(used) / 1024, float(total) / 1024  # 转GB

used_mem, total_mem = get_gpu_memory()
print(f"\n--- 模型加载后显存（nvidia-smi） ---")
print(f"已使用：{used_mem:.2f} GB / 总计：{total_mem:.2f} GB | 剩余：{total_mem - used_mem:.2f} GB")

# ===================== 4. FastAPI Chat接口 =====================
app = FastAPI(title="Qwen3-32B Chat API", description="OpenAI兼容Chat接口")

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ===================== 高血压风险筛查提示词模板 =====================
HTN_SCREENING_PROMPT_TEMPLATE = """依据《中国高血压防治指南2023》《中国2型糖尿病防治指南2020》《慢性肾脏病筛查防治指南2022》，基于患者客观数据逐条比对，禁止虚构。

筛查COUNT项高血压易患因素：
FACTOR_LINES

match=1:有明确数据证实;match=0:未提及/否定/数据正常
reason须简短(≤30字)。仅返回JSON，无多余文字：
[JSON_EXAMPLE]

患者信息：PATIENT_DATA_PLACEHOLDER"""


def build_screening_prompt(factor_list, patient_data):
    """根据外部传入的因素编码和名称动态构建提示词"""
    count = len(factor_list)
    factor_lines = "\n".join(
        f"{i+1}.{f['name']}" for i, f in enumerate(factor_list)
    )
    # 用第一个因素构造JSON示例，i字段用编码
    json_example = json.dumps(
        {"i": factor_list[0]["code"], "t": factor_list[0]["name"], "m": 0, "r": "简要原因"},
        ensure_ascii=False
    )
    prompt = HTN_SCREENING_PROMPT_TEMPLATE.replace("COUNT", str(count))
    prompt = prompt.replace("FACTOR_LINES", factor_lines)
    prompt = prompt.replace("JSON_EXAMPLE", json_example)
    prompt = prompt.replace("PATIENT_DATA_PLACEHOLDER", patient_data)
    return prompt

# 请求/响应模型
class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str

class ChatRequest(BaseModel):
    model: str = "qwen3-32b"
    messages: List[ChatMessage]
    temperature: float = Field(default=0.7, ge=0, le=2)
    top_p: float = Field(default=0.8, ge=0, le=1)
    max_tokens: int = Field(default=512, ge=1, le=1024)
    stream: bool = False  # 暂不支持流式，预留字段

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    """OpenAI兼容的Chat接口"""
    # 格式化对话（messages → chat_template）
    # 追加 /no_think 关闭Qwen3思考模式，避免输出大段<think...</think)>
    msg_dicts = [{"role": m.role, "content": m.content} for m in req.messages]
    if msg_dicts and msg_dicts[-1]["role"] == "user":
        msg_dicts[-1]["content"] += " /no_think"
    prompt = tokenizer.apply_chat_template(msg_dicts, tokenize=False, add_generation_prompt=True)

    sampling_params = SamplingParams(
        temperature=req.temperature,
        top_p=req.top_p,
        max_tokens=req.max_tokens,
    )

    start_time = time.time()
    outputs = llm.generate([prompt], sampling_params)
    elapsed = time.time() - start_time

    completion = outputs[0].outputs[0]
    # 过滤残留的<think...</think)>标签（双重保险）
    text = re.sub(r"<think.*?</think\s*>", "", completion.text, flags=re.DOTALL).strip()
    prompt_tokens = len(outputs[0].prompt_token_ids)
    new_tokens = len(completion.token_ids)

    # OpenAI兼容响应格式
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": req.model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": text},
            "finish_reason": "stop",
        }],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": new_tokens,
            "total_tokens": prompt_tokens + new_tokens,
        },
        "speed_tokens_per_sec": round(new_tokens / elapsed, 2) if elapsed > 0 else 0,
    }

@app.get("/v1/models")
async def list_models():
    """模型列表接口"""
    return {
        "object": "list",
        "data": [{"id": "qwen3-32b", "object": "model", "owned_by": "local"}]
    }

@app.get("/health")
async def health():
    """健康检查接口"""
    used, total = get_gpu_memory()
    return {"status": "ok", "gpu_used_gb": round(used, 2), "gpu_total_gb": round(total, 2)}

# ===================== 高血压易患因素筛查接口 =====================
# 默认因素列表（当外部未传入factorList时使用）
DEFAULT_FACTORS = [
    {"code": "HTN_F01", "name": "正常高值血压(SBP120-139/DBP80-89mmHg)"},
    {"code": "HTN_F02", "name": "年龄≥45岁"},
    {"code": "HTN_F03", "name": "超重(BMI≥24)或腹型肥胖(男腰围≥90/女≥85cm)"},
    {"code": "HTN_F04", "name": "高血压家族史(一二级亲属)"},
    {"code": "HTN_F05", "name": "高盐饮食(日盐>6g)"},
    {"code": "HTN_F06", "name": "长期大量饮酒(男日乙醇≥40g/女≥20g)"},
    {"code": "HTN_F07", "name": "吸烟(含被动)"},
    {"code": "HTN_F08", "name": "缺乏体力劳动(周运动<150min/久坐)"},
    {"code": "HTN_F09", "name": "长期精神紧张"},
    {"code": "HTN_F10", "name": "血脂异常(TC≥5.18/LDL-C≥3.37/TG≥1.70mmol/L或HDL-C降低)"},
    {"code": "HTN_F11", "name": "糖尿病(FPG≥7.0/HbA1c≥6.5%/已诊断)"},
]

class FactorItem(BaseModel):
    code: str  # 因素编码
    name: str  # 因素名称

class PatientDataRequest(BaseModel):
    chatContent: str
    factorList: List[FactorItem] = []  # 外部传入的因素编码和名称，为空时使用默认列表

@app.post("/api/htn-screening")
async def htn_screening(req: PatientDataRequest):
    """高血压易患因素筛查接口"""
    # 使用外部传入的因素列表，若为空则使用默认列表
    factor_list = [f.model_dump() for f in req.factorList] if req.factorList else DEFAULT_FACTORS

    # 构建编码→名称映射（用于结果替换）
    code_map = {f["code"]: f["name"] for f in factor_list}

    # 动态构建提示词
    prompt = build_screening_prompt(factor_list, req.chatContent)
    messages = [
        {"role": "user", "content": prompt + " /no_think"},
    ]
    formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    # 计算剩余可用token：总上下文 - 输入prompt = 最大输出
    prompt_token_ids = tokenizer.encode(formatted)
    prompt_len = len(prompt_token_ids)
    max_output = max(64, 4096 - prompt_len)  # 至少保留64 tokens输出

    sampling_params = SamplingParams(temperature=0.1, top_p=0.8, max_tokens=max_output)
    outputs = llm.generate([formatted], sampling_params)

    completion = outputs[0].outputs[0]
    text = re.sub(r"<think.*?</think\s*>", "", completion.text, flags=re.DOTALL).strip()
    # 去掉markdown代码块包裹
    text = re.sub(r"^```json\s*", "", text)
    text = re.sub(r"\s*```$", "", text)

    # 截断修复：如果输出被max_tokens截断，JSON不完整，尝试补全
    if completion.finish_reason == "length" and "[" in text and "]" not in text:
        # 去掉最后一个不完整的对象，补上闭合括号
        last_brace = text.rfind("},")
        if last_brace > 0:
            text = text[:last_brace + 1] + "\n]"
        else:
            last_brace = text.rfind("{")
            if last_brace > 0:
                text = text[:last_brace] + "]"

    try:
        json_match = re.search(r'\[.*\]', text, re.DOTALL)
        if json_match:
            screening_list = json.loads(json_match.group())
        else:
            screening_list = json.loads(text)
        # 构建序号→编码映射（LLM可能返回数字序号或字符串序号1,2,...而非编码）
        idx_to_code = {i + 1: f["code"] for i, f in enumerate(factor_list)}
        # 短字段名映射为完整字段名（逐条处理，单条异常不影响其他）
        result_list = []
        for item in screening_list:
            try:
                raw_i = item.get("i", "")
                try:
                    raw_i_int = int(raw_i)
                except (ValueError, TypeError):
                    raw_i_int = None
                if raw_i_int is not None and raw_i_int in idx_to_code:
                    code = idx_to_code[raw_i_int]
                elif raw_i in code_map:
                    code = raw_i
                else:
                    code = str(raw_i)
                match_val = item.get("m", 0)
                result_list.append({
                    "code": code,
                    "name": item.get("t", code_map.get(code, "")),
                    "match": match_val,
                    "reason": item.get("r", ""),
                })
            except Exception as e:
                print(f"单条解析异常：{e}, item={item}")
                result_list.append(item)

        return {"status": 200, "data": result_list, "truncated": completion.finish_reason == "length"}
    except json.JSONDecodeError as e:
        print(f"JSON解析异常：{e}\n原始返回：{text}")
        return {"status": 500, "data": [], "error": f"模型返回格式异常：{str(e)}", "raw": text}
    except Exception as e:
        print(f"接口异常：{e}")
        return {"status": 500, "data": [], "error": str(e)}

# ===================== 5. 启动 =====================
if __name__ == "__main__":
    print("\n模型加载完成！启动Chat API服务...")
    print(f"Chat接口：POST http://0.0.0.0:8000/v1/chat/completions")
    print(f"高血压筛查：POST http://0.0.0.0:8000/api/htn-screening")
    print(f"模型列表：GET  http://0.0.0.0:8000/v1/models")
    print(f"健康检查：GET  http://0.0.0.0:8000/health")
    uvicorn.run(app, host="0.0.0.0", port=8000)

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModel, AutoConfig
from pydantic import BaseModel, Field
from typing import List, Literal
import torch
import uvicorn
from modelscope import snapshot_download
import time
import uuid
import re
import json
import warnings
warnings.filterwarnings("ignore")

# --- 全局变量 ---
model = None
tokenizer = None

model_name = "ZhipuAI/chatglm3-6b"
cache_dir = "/home/model"

# 2. 初始化FastAPI应用
app = FastAPI(title="ChatGLM3-6B Chat API", description="OpenAI兼容Chat接口 + 高血压筛查")

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 3. 加载模型和Tokenizer
print("正在下载/校验模型缓存...")
local_model_path = snapshot_download(model_name, cache_dir=cache_dir)

print(f"正在加载模型，路径: {local_model_path}")
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)

# 修复ChatGLM配置兼容性问题
config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=True)
if not hasattr(config, 'max_length'):
    config.max_length = config.seq_length if hasattr(config, 'seq_length') else 8192

# 使用AutoModel加载
try:
    if not torch.cuda.is_available():
        raise RuntimeError("未检测到GPU，请确保CUDA环境配置正确")

    print(f"检测到GPU: {torch.cuda.get_device_name(0)}")
    print(f"GPU显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")

    model = AutoModel.from_pretrained(
        local_model_path,
        config=config,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    model.eval()

    device = next(model.parameters()).device
    print(f"模型加载成功！当前设备: {device}")

    # 显存统计
    alloc_mem = torch.cuda.memory_allocated() / 1024**3
    reserv_mem = torch.cuda.memory_reserved() / 1024**3
    total_mem = torch.cuda.get_device_properties(0).total_memory / 1024**3
    print(f"已分配：{alloc_mem:.2f} GB | 已保留：{reserv_mem:.2f} GB | 剩余：{total_mem - reserv_mem:.2f} GB / 总计：{total_mem:.2f} GB")

except Exception as e:
    print(f"加载失败: {e}")
    model = None

# ===================== 高血压风险筛查提示词模板 =====================
HTN_SCREENING_PROMPT_TEMPLATE = """依据《中国高血压防治指南2023》《中国2型糖尿病防治指南2020》《慢性肾脏病筛查防治指南2022》，基于患者客观数据逐条比对，禁止虚构。

筛查{count}项高血压易患因素：
{factor_lines}

match=1:有明确数据证实;match=0:未提及/否定/数据正常
reason须简短(≤30字)。仅返回JSON，无多余文字：
[{json_example}]

患者信息：PATIENT_DATA_PLACEHOLDER"""


def build_screening_prompt(factor_list, patient_data):
    """根据外部传入的因素编码和名称动态构建提示词"""
    count = len(factor_list)
    factor_lines = "\n".join(
        f"{i+1}.{f['name']}" for i, f in enumerate(factor_list)
    )
    # 用第一个因素构造JSON示例，i字段用编码
    json_example = json.dumps(
        {"i": factor_list[0]["code"], "t": factor_list[0]["name"], "m": 0, "r": "简要原因"},
        ensure_ascii=False
    )
    prompt = HTN_SCREENING_PROMPT_TEMPLATE.replace("{count}", str(count))
    prompt = prompt.replace("{factor_lines}", factor_lines)
    prompt = prompt.replace("{json_example}", json_example)
    prompt = prompt.replace("PATIENT_DATA_PLACEHOLDER", patient_data)
    return prompt

# ===================== 请求/响应模型 =====================
class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str

class ChatRequest(BaseModel):
    model: str = "chatglm3-6b"
    messages: List[ChatMessage]
    temperature: float = Field(default=0.7, ge=0, le=2)
    top_p: float = Field(default=0.8, ge=0, le=1)
    max_tokens: int = Field(default=512, ge=1, le=2048)
    stream: bool = False

class FactorItem(BaseModel):
    code: str  # 因素编码
    name: str  # 因素名称

class PatientDataRequest(BaseModel):
    chatContent: str
    factorList: List[FactorItem] = []  # 外部传入的因素编码和名称，为空时使用默认列表

# ===================== Chat对话接口 =====================
@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    """OpenAI兼容的Chat接口"""
    if model is None:
        raise HTTPException(status_code=500, detail="模型未加载，请检查服务器启动日志")

    # 拼接对话历史为ChatGLM格式
    history = []
    user_query = ""
    for m in req.messages:
        if m.role == "user":
            user_query = m.content
        elif m.role == "assistant" and user_query:
            history.append((user_query, m.content))
            user_query = ""

    # 取最后一条user消息作为当前输入
    if not user_query:
        user_query = req.messages[-1].content

    start_time = time.time()
    try:
        response, _ = model.chat(tokenizer, user_query, history=history)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"推理出错: {str(e)}")
    elapsed = time.time() - start_time

    # 粗略token统计（ChatGLM无精确prompt_token_ids）
    prompt_tokens = len(tokenizer.encode(user_query))
    completion_tokens = len(tokenizer.encode(response))
    speed = completion_tokens / elapsed if elapsed > 0 else 0

    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": req.model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": response},
            "finish_reason": "stop",
        }],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
        "speed_tokens_per_sec": round(speed, 2),
    }

# ===================== 模型列表 =====================
@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{"id": "chatglm3-6b", "object": "model", "owned_by": "local"}]
    }

# ===================== 健康检查 =====================
@app.get("/health")
async def health():
    if torch.cuda.is_available():
        alloc_mem = torch.cuda.memory_allocated() / 1024**3
        total_mem = torch.cuda.get_device_properties(0).total_memory / 1024**3
        return {"status": "ok", "gpu_used_gb": round(alloc_mem, 2), "gpu_total_gb": round(total_mem, 2)}
    return {"status": "ok"}

# ===================== 高血压易患因素筛查接口 =====================
# 默认因素列表（当外部未传入factorList时使用）
DEFAULT_FACTORS = [
    {"code": "HTN_F01", "name": "正常高值血压(SBP120-139/DBP80-89mmHg)"},
    {"code": "HTN_F02", "name": "年龄≥45岁"},
    {"code": "HTN_F03", "name": "超重(BMI≥24)或腹型肥胖(男腰围≥90/女≥85cm)"},
    {"code": "HTN_F04", "name": "高血压家族史(一二级亲属)"},
    {"code": "HTN_F05", "name": "高盐饮食(日盐>6g)"},
    {"code": "HTN_F06", "name": "长期大量饮酒(男日乙醇≥40g/女≥20g)"},
    {"code": "HTN_F07", "name": "吸烟(含被动)"},
    {"code": "HTN_F08", "name": "缺乏体力劳动(周运动<150min/久坐)"},
    {"code": "HTN_F09", "name": "长期精神紧张"},
    {"code": "HTN_F10", "name": "血脂异常(TC≥5.18/LDL-C≥3.37/TG≥1.70mmol/L或HDL-C降低)"},
    {"code": "HTN_F11", "name": "糖尿病(FPG≥7.0/HbA1c≥6.5%/已诊断)"},
]

@app.post("/api/htn-screening")
async def htn_screening(req: PatientDataRequest):
    """高血压易患因素筛查接口"""
    if model is None:
        raise HTTPException(status_code=500, detail="模型未加载，请检查服务器启动日志")

    # 使用外部传入的因素列表，若为空则使用默认列表
    factor_list = [f.model_dump() for f in req.factorList] if req.factorList else DEFAULT_FACTORS

    # 构建编码→名称映射（用于结果替换）
    code_map = {f["code"]: f["name"] for f in factor_list}

    # 动态构建提示词
    prompt = build_screening_prompt(factor_list, req.chatContent)

    # ChatGLM3用model.chat推理
    start_time = time.time()
    try:
        response, _ = model.chat(
            tokenizer,
            prompt,
            history=[],
        )
    except Exception as e:
        return {"code": 500, "screeningList": [], "error": f"推理出错: {str(e)}"}
    elapsed = time.time() - start_time

    # 过滤残留标签
    text = re.sub(r"<think.*?</think\s*>", "", response, flags=re.DOTALL).strip()
    # 去掉markdown代码块包裹
    text = re.sub(r"^```json\s*", "", text)
    text = re.sub(r"\s*```$", "", text)

    # 截断修复：JSON不完整时尝试补全
    if "[" in text and "]" not in text:
        last_brace = text.rfind("},")
        if last_brace > 0:
            text = text[:last_brace + 1] + "\n]"
        else:
            last_brace = text.rfind("{")
            if last_brace > 0:
                text = text[:last_brace] + "]"

    try:
        json_match = re.search(r'\[.*\]', text, re.DOTALL)
        if json_match:
            screening_list = json.loads(json_match.group())
        else:
            screening_list = json.loads(text)
        # 构建序号→编码映射（LLM可能返回数字序号1,2,...而非编码）
        idx_to_code = {i + 1: f["code"] for i, f in enumerate(factor_list)}
        # 短字段名映射为完整字段名
        matched_codes = set()
        for item in screening_list:
            raw_i = item.pop("i", "")
            # 如果是数字序号，映射为对应编码；如果已经是编码则直接使用
            if isinstance(raw_i, int) and raw_i in idx_to_code:
                code = idx_to_code[raw_i]
            elif raw_i in code_map:
                code = raw_i
            else:
                code = str(raw_i)
            item["code"] = code
            item["title"] = item.pop("t", code_map.get(code, ""))
            item["match"] = item.pop("m", 0)
            item["reason"] = item.pop("r", "")
            matched_codes.add(code)

        # 补全：LLM未返回的因素自动填充 match=0
        for f in factor_list:
            if f["code"] not in matched_codes:
                screening_list.append({
                    "code": f["code"],
                    "title": f["name"],
                    "match": 0,
                    "reason": "未检出相关数据",
                })

        return {"code": 200, "screeningList": screening_list, "truncated": False}
    except json.JSONDecodeError as e:
        print(f"JSON解析异常：{e}\n原始返回：{text}")
        return {"code": 500, "screeningList": [], "error": f"模型返回格式异常：{str(e)}", "raw": text}
    except Exception as e:
        print(f"接口异常：{e}")
        return {"code": 500, "screeningList": [], "error": str(e)}

# ===================== 启动 =====================
if __name__ == "__main__":
    print("\n模型加载完成！启动Chat API服务...")
    print(f"Chat接口：POST http://0.0.0.0:8000/v1/chat/completions")
    print(f"高血压筛查：POST http://0.0.0.0:8000/api/htn-screening")
    print(f"模型列表：GET  http://0.0.0.0:8000/v1/models")
    print(f"健康检查：GET  http://0.0.0.0:8000/health")
    uvicorn.run(app, host="0.0.0.0", port=8000)