diff --git a/docker-compose.yml b/docker-compose.yml index 9b4ab67..311313d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -3,11 +3,11 @@ services: build: . container_name: llm-compass ports: - - "8000:8000" + - "402:8000" environment: - DASHSCOPE_API_KEY=${DASHSCOPE_API_KEY} volumes: - - compass-data:/app/data + - ./data:/app/data restart: unless-stopped healthcheck: test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"] diff --git a/nvidia_router.py b/nvidia_router.py index 7548a84..382385d 100644 --- a/nvidia_router.py +++ b/nvidia_router.py @@ -10,11 +10,12 @@ NVIDIA Prompt Task & Complexity Classifier Router import torch import torch.nn as nn -from transformers import AutoTokenizer, DebertaV2Model, AutoConfig +from transformers import AutoTokenizer, DebertaV2Model, DebertaV2Config from safetensors.torch import load_file from huggingface_hub import hf_hub_download from typing import Dict, Optional import logging +import json logger = logging.getLogger(__name__) @@ -115,7 +116,16 @@ class NvidiaComplexityRouter: # Creativity 映射 CREATIVITY_MAP = {0: "High", 1: "Low", 2: "No"} - def __init__(self, device: str = "cpu"): + def __init__(self, device: str = "auto"): + if device == "auto": + if torch.backends.mps.is_available(): + device = "mps" + logger.info("MPS (Metal GPU) detected, using MPS acceleration") + elif torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + logger.info("No GPU detected, using CPU") self.device = device self.tokenizer = None self.model = None @@ -129,8 +139,18 @@ class NvidiaComplexityRouter: logger.info(f"Loading NVIDIA classifier: {self.MODEL_NAME}") - # 1. 加载 config - self.config = AutoConfig.from_pretrained(self.MODEL_NAME) + # 1. 手动加载自定义 config.json(该模型无 model_type,AutoConfig 不兼容) + config_path = hf_hub_download(self.MODEL_NAME, "config.json") + with open(config_path, "r") as f: + custom_config = json.load(f) + + # 构建 backbone 的 DeBERTa config(从 base_model 加载) + base_model = custom_config.get("base_model", "microsoft/DeBERTa-v3-base") + self.config = DebertaV2Config.from_pretrained(base_model) + # 保存自定义分类头参数 + self.config.target_sizes = custom_config["target_sizes"] + self.config.fc_dropout = custom_config.get("fc_dropout", 0.2) + self.config.base_model = base_model # 2. 加载 tokenizer (slow模式,兼容性好) self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME, use_fast=False) @@ -143,9 +163,13 @@ class NvidiaComplexityRouter: self.model.load_state_dict(state_dict, strict=False) self.model.to(self.device) + # MPS 需要 float16 以避免矩阵乘法数据类型冲突 + if self.device == "mps": + self.model.half() self.model.eval() self._initialized = True - logger.info("NVIDIA classifier loaded successfully") + dtype = "float16" if self.device == "mps" else "float32" + logger.info(f"NVIDIA classifier loaded successfully on {self.device} ({dtype})") def predict(self, query: str) -> Dict: """