初始化提交

yuanqi/src/engine/asr.py

@@ -0,0 +1,298 @@
+"""
+本地 ASR 推理引擎 — FunASR + SenseVoiceSmall。
+
+基于阿里达摩院 FunASR 框架加载 SenseVoiceSmall 模型，提供本地离线语音转文字能力。
+
+支持语言：中文、粤语、英文、日文、韩文等。
+额外能力：情感识别、事件检测（笑声、掌声、音乐等）。
+
+模型下载位置（按优先级）：
+1. ``YUANQI_ASR_MODEL`` 为本地路径时，直接从该路径加载
+2. ``YUANQI_ASR_MODEL`` 为 HuggingFace/ModelScope 模型名时，下载到
+   ``models/funasr/<模型名>`` 后再加载
+"""
+
+from __future__ import annotations
+
+import logging as _logging
+import os
+import re
+import tempfile
+from pathlib import Path
+from typing import Any
+
+# 在 funasr / modelscope 导入前压制其冗余 INFO 日志
+_logging.getLogger("modelscope").setLevel(_logging.WARNING)
+_logging.getLogger("modelscope_hub").setLevel(_logging.WARNING)
+
+from src.config import settings  # noqa: E402
+from src.utils.logging import logger  # noqa: E402
+
+# funasr 支持的音频格式白名单
+_SUPPORTED_FORMATS = frozenset({
+    "wav", "mp3", "flac", "ogg", "m4a", "opus",
+    "aac", "pcm", "amr", "webm",
+})
+
+# 语言代码映射（ISO-639-1 → SenseVoice 内部标签）
+_LANG_MAP: dict[str, str] = {
+    "auto": "auto",   # 自动检测
+    "zh": "zh",       # 中文
+    "en": "en",       # 英文
+    "ja": "ja",       # 日文
+    "ko": "ko",       # 韩文
+    "yue": "yue",     # 粤语
+    "pt": "pt",       # 葡萄牙语
+    "fr": "fr",       # 法语
+    "es": "es",       # 西班牙语
+    "ru": "ru",       # 俄语
+    "ar": "ar",       # 阿拉伯语
+    "vi": "vi",       # 越南语
+    "th": "th",       # 泰语
+    "id": "id",       # 印度尼西亚语
+}
+
+
+# SenseVoiceSmall 情感标签
+_SENSEVOICE_EMOTIONS = frozenset({
+    "NEUTRAL", "HAPPY", "SAD", "ANGRY", "SURPRISED",
+    "FEARFUL", "DISGUSTED", "EMO_UNKNOWN",
+})
+
+# SenseVoiceSmall 事件标签
+_SENSEVOICE_EVENTS = frozenset({
+    "Speech", "laughter", "applause", "music", "sing",
+    "cry", "breath", "cough", "sneeze", "nospeech",
+})
+
+# SenseVoiceSmall 内部标志（不暴露给用户）
+_SENSEVOICE_INTERNAL = frozenset({
+    "withitn", "woitn",
+})
+
+# 匹配所有 <|...|> 标签
+_TAG_PATTERN = re.compile(r"<\|([^|]+)\|>")
+
+
+def _clean_sensevoice_text(raw: str) -> tuple[str, str | None, list[str], list[str]]:
+    """清洗 SenseVoiceSmall 原始输出，提取标签并从文本中剥离。
+
+    SenseVoiceSmall 返回示例::
+
+        <|zh|><|NEUTRAL|><|Speech|><|withitn|>你好呀，今天天气真好。
+
+    Returns:
+        (cleaned_text, detected_lang, emotions, events)
+    """
+    tags = _TAG_PATTERN.findall(raw)
+    cleaned = _TAG_PATTERN.sub("", raw).strip()
+
+    detected_lang = None
+    emotions: list[str] = []
+    events: list[str] = []
+
+    for tag in tags:
+        t = tag.strip()
+        if t in _SENSEVOICE_EMOTIONS:
+            emotions.append(t)
+        elif t in _SENSEVOICE_EVENTS:
+            events.append(t)
+        elif t in _SENSEVOICE_INTERNAL:
+            pass  # 内部标志不返回
+        elif t in _LANG_MAP:
+            detected_lang = t
+        # 其他未知标签忽略
+
+    return cleaned, detected_lang, emotions, events
+
+
+class SenseVoiceASREngine:
+    """本地 ASR 引擎 — FunASR + SenseVoiceSmall。"""
+
+    name: str = "asr"
+    model_loaded: bool = False
+
+    def __init__(self) -> None:
+        self._model_name = settings.asr_model or "iic/SenseVoiceSmall"
+        self._model = None
+        logger.info(
+            "ASR 引擎配置 | model={}",
+            self._model_name,
+        )
+
+    def _resolve_model_path(self) -> str:
+        """解析模型路径并下载到本地。
+
+        同 embedding / reranker 的路径解析模式。
+        """
+        # 1) 直接是本地路径
+        direct = Path(self._model_name)
+        if direct.is_absolute() and direct.exists():
+            return str(direct)
+        candidate = settings.models_dir_resolved / self._model_name
+        if candidate.exists():
+            return str(candidate)
+
+        # 2) 已下载到 models/funasr/<repo_name>
+        repo_name = Path(self._model_name).name
+        target_dir = settings.models_dir_resolved / "funasr" / repo_name
+        if target_dir.exists():
+            return str(target_dir)
+
+        # 3) 下载到 models/funasr/<repo_name>
+        target = target_dir
+        logger.info("本地 ASR 模型不存在，从 {} 下载到 {}", settings.model_source, target)
+        from src.utils.downloader import download_model
+
+        if download_model(self._model_name, target, settings.model_source):
+            return str(target)
+        logger.warning("ASR 模型下载失败，回退到 HF 缓存加载")
+        return self._model_name
+
+    async def health(self) -> dict[str, Any]:
+        if not self.model_loaded:
+            return {"status": "unavailable", "model": self._model_name}
+        return {"status": "ok", "model": self._model_name}
+
+    async def load(self) -> None:
+        """加载 FunASR 模型（SenseVoiceSmall）。
+
+        模型加载完成后，引擎即可使用 ``transcribe()`` 对音频进行转写。
+        """
+        if self.model_loaded:
+            return
+        try:
+            model_path = self._resolve_model_path()
+            from funasr import AutoModel
+
+            logger.info("正在加载 ASR 模型: {}", model_path)
+            self._model = AutoModel(
+                model="iic/SenseVoiceSmall",
+                model_dir=model_path,
+                disable_update=True,
+                log_level="WARNING",
+            )
+            self.model_loaded = True
+            logger.info("ASR 模型加载完成: {}", model_path)
+        except ImportError:
+            logger.error(
+                "funasr 未安装。请运行: pip install 'hunjing-yuanqi[asr]' 或 "
+                "cd yuanqi && uv add funasr"
+            )
+        except Exception as e:
+            logger.error("ASR 模型加载失败: {}", e)
+
+    def _normalize_language(self, language: str | None) -> str:
+        """将用户指定的语言代码归一化为 SenseVoice 支持的标签。
+
+        完整支持语言列表见 FunASR 文档。
+        """
+        if not language or language.lower() in ("auto", ""):
+            return "auto"
+        lang = _LANG_MAP.get(language.lower())
+        return lang if lang else "auto"
+
+    async def transcribe(
+        self,
+        audio_data: bytes,
+        format: str = "wav",
+        language: str | None = None,
+    ) -> dict[str, Any]:
+        """将音频字节转写为文本。
+
+        Args:
+            audio_data: 原始音频字节。
+            format: 音频容器格式（如 ``wav``、``mp3``、``m4a``）。
+            language: 可选的语言提示（ISO-639-1），不传则自动检测。
+
+        Returns:
+            dict:
+                - ``text``: 转写文本
+                - ``model``: 使用的模型名
+                - ``language``: 检测到的语言
+                - ``emotions``: 可选的情感标签（如果模型返回）
+                - ``events``: 可选的事件标签（笑声/掌声等）
+        """
+        if not self.model_loaded or self._model is None:
+            logger.warning("ASR 模型未加载，返回空结果")
+            return {"text": "", "model": self._model_name, "language": None}
+
+        lang_tag = self._normalize_language(language)
+
+        # 将音频字节写入临时文件进行处理
+        suffix = f".{format}" if format.startswith(".") else f".{format}"
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+            tmp.write(audio_data)
+            tmp_path = tmp.name
+
+        try:
+            result = self._model.generate(
+                input=tmp_path,
+                language=lang_tag,
+                use_itn=True,   # 逆文本正规化（数字、标点、单位）
+                ban_emo_unk=False,  # 保留情感/事件标签
+                # 可选：hotword="xxx", 热词增强
+            )
+
+            # 解析结果 —— 格式：[{"key": "xxx", "text": "...", "timestamp": "..."}]
+            text = ""
+            detected_lang = language
+            emotions: list[str] = []
+            events: list[str] = []
+
+            if isinstance(result, list) and len(result) > 0:
+                raw_text = (
+                    result[0].get("text", "")
+                    if isinstance(result[0], dict)
+                    else str(result[0])
+                )
+
+                # 清洗 SenseVoiceSmall 内部标签，提取元信息
+                cleaned_text, tag_lang, tag_emotions, tag_events = _clean_sensevoice_text(raw_text)
+                text = cleaned_text
+
+                # 模型检测到的语言优先于用户指定的 language hint
+                if tag_lang and language in (None, "auto"):
+                    detected_lang = tag_lang
+
+                emotions = tag_emotions
+                events = tag_events
+
+                # 也检查 result dict 中的额外字段
+                if isinstance(result[0], dict):
+                    emo = result[0].get("emotion")
+                    if emo and emo not in emotions:
+                        emotions.append(emo)
+
+            logger.info(
+                "ASR 转写完成 | format={} lang={} text_len={}",
+                format, detected_lang, len(text),
+            )
+            if text:
+                logger.info("ASR 识别结果: {}", text)
+            return {
+                "text": text,
+                "model": self._model_name,
+                "language": detected_lang or "auto",
+                "emotions": emotions if emotions else None,
+                "events": events if events else None,
+            }
+
+        except Exception as e:
+            logger.error("ASR 转写失败: {}", e)
+            return {"text": "", "model": self._model_name, "language": None}
+
+        finally:
+            # 清理临时文件
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+
+    async def unload(self) -> None:
+        self._model = None
+        self.model_loaded = False
+        logger.info("ASR 模型已卸载")
+
+    async def close(self) -> None:
+        await self.unload()

yuanqi/src/engine/base.py

@@ -0,0 +1,36 @@
+"""
+推理引擎抽象基类。
+
+定义元炁支持的所有引擎的统一接口。所有引擎独立加载模型、独立管理生命周期。
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class InferenceEngine(ABC):
+    """推理引擎基类。"""
+
+    name: str = "base"
+    model_loaded: bool = False
+
+    @abstractmethod
+    async def health(self) -> dict[str, Any]:
+        """返回引擎健康状态。"""
+        ...
+
+    @abstractmethod
+    async def load(self) -> None:
+        """加载模型（启动时或首次调用时懒加载）。"""
+        ...
+
+    @abstractmethod
+    async def unload(self) -> None:
+        """卸载模型，释放资源。"""
+        ...
+
+    async def close(self) -> None:
+        """释放引擎资源。"""
+        await self.unload()

yuanqi/src/engine/embedding.py

@@ -0,0 +1,197 @@
+"""
+本地 Embedding 推理引擎。
+
+基于 FlagEmbedding（BAAI 官方库）加载 BGE 系列嵌入模型（如 BGE-M3），
+支持稠密（dense）、稀疏（sparse）和多向量（colbert）三种嵌入模式，
+默认输出稠密向量（dense embedding）。
+
+模型下载位置（按优先级）：
+1. ``YUANQI_EMBEDDING_MODEL`` 为本地路径时，直接从该路径加载
+2. ``YUANQI_EMBEDDING_MODEL`` 为 HuggingFace 模型名时，下载到
+   ``models/FlagModel/<模型名>`` 后再加载
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from src.config import settings
+from src.utils.device import detect_device
+from src.utils.logging import logger
+
+
+class EmbeddingEngine:
+    """本地 Embedding 引擎。"""
+
+    name: str = "embedding"
+    model_loaded: bool = False
+
+    def __init__(self) -> None:
+        self._device = detect_device(settings.device)
+        self._model_name = settings.embedding_model or "FlagModel/bge-m3"
+        self._model = None
+        self._tokenizer = None  # 用于 token 计数
+        logger.info(
+            "Embedding 引擎配置 | model={} device={}",
+            self._model_name,
+            self._device,
+        )
+
+    def _resolve_model_path(self) -> str:
+        """解析模型路径。
+
+        ``embedding_model`` 是 HuggingFace 模型名时，下载到
+        ``models/FlagModel/<模型名>`` 后再加载。
+        """
+        # 1) 直接是本地路径
+        direct = Path(self._model_name)
+        if direct.is_absolute() and direct.exists():
+            return str(direct)
+        candidate = settings.models_dir_resolved / self._model_name
+        if candidate.exists():
+            return str(candidate)
+
+        # 2) 已下载到目标目录
+        target_dir = settings.models_dir_resolved / "FlagModel" / self._model_name
+        if target_dir.exists():
+            return str(target_dir)
+
+        # 3) 已下载到 models/FlagModel/<repo_name>
+        repo_name = Path(self._model_name).name
+        target_dir2 = settings.models_dir_resolved / "FlagModel" / repo_name
+        if target_dir2.exists():
+            return str(target_dir2)
+
+        # 4) 下载到 models/FlagModel/<repo_name>
+        target = target_dir2
+        logger.info("本地模型不存在，从 {} 下载到 {}", settings.model_source, target)
+        from src.utils.downloader import download_model
+
+        if download_model(self._model_name, target, settings.model_source):
+            return str(target)
+        logger.warning("模型下载失败，回退到 HF 缓存加载")
+        return self._model_name
+
+    async def health(self) -> dict[str, Any]:
+        if not self.model_loaded:
+            return {"status": "unavailable", "model": self._model_name}
+        return {"status": "ok", "model": self._model_name}
+
+    async def load(self) -> None:
+        """加载 FlagEmbedding 模型和对应的 tokenizer。"""
+        if self.model_loaded:
+            return
+        try:
+            model_path = self._resolve_model_path()
+            from FlagEmbedding import BGEM3FlagModel
+
+            logger.info("正在加载 Embedding 模型: {}", model_path)
+            self._model = BGEM3FlagModel(
+                model_path,
+                use_fp16=self._device in ("cuda", "mps"),
+                device=self._device,
+            )
+            # 加载同 tokenizer 用于计数
+            try:
+                from transformers import AutoTokenizer
+                self._tokenizer = AutoTokenizer.from_pretrained(model_path)
+            except Exception:
+                self._tokenizer = None
+            self.model_loaded = True
+            logger.info("Embedding 模型加载完成: {}", model_path)
+        except ImportError:
+            logger.error(
+                "FlagEmbedding 未安装。请运行: pip install 'hunjing-yuanqi[embedding]'"
+            )
+        except Exception as e:
+            logger.error("Embedding 模型加载失败: {}", e)
+
+    def count_tokens(self, texts: list[str]) -> list[int]:
+        """统计每段文本的 token 数。"""
+        from src.utils.token_counter import TokenCounter, EmbeddingTokenCountStrategy
+
+        if self._tokenizer is None:
+            return [0] * len(texts)
+        strategy = EmbeddingTokenCountStrategy()
+        return [
+            strategy.count(self._tokenizer, text=t) for t in texts
+        ]
+
+    async def embed(
+        self,
+        texts: list[str],
+        encoding_format: str = "dense",
+    ) -> list[dict[str, Any]]:
+        """将文本列表转为向量嵌入。
+
+        Args:
+            texts: 待编码的文本列表。
+            encoding_format: 嵌入格式。
+                ``dense`` — 稠密向量（默认，兼容 OpenAI API）。
+                ``sparse`` — 稀疏向量 {indices, values}。
+                ``colbert`` — ColBERT 多向量。
+                ``all`` — 同时返回 dense + sparse + colbert。
+
+        Returns:
+            向量列表，每项为 dict，不同格式含不同 key：
+
+            - ``dense``: ``{"embedding": [float, ...]}``
+            - ``sparse``: ``{"sparse": {"indices": [int, ...], "values": [float, ...]}}``
+            - ``colbert``: ``{"colbert": [[float, ...], ...]}``
+            - ``all``: 同时含以上三个 key。
+        """
+        if not self.model_loaded or self._model is None:
+            empty: list[dict[str, Any]] = [{"embedding": []} for _ in texts]
+            return empty
+
+        return_sparse = encoding_format in ("sparse", "all")
+        return_colbert = encoding_format in ("colbert", "all")
+
+        output = self._model.encode(
+            texts,
+            return_dense=True,
+            return_sparse=return_sparse,
+            return_colbert_vecs=return_colbert,
+        )
+
+        results: list[dict[str, Any]] = []
+
+        for i in range(len(texts)):
+            item: dict[str, Any] = {}
+
+            # 稠密
+            dense_vec = output["dense_vecs"][i].tolist()
+            item["embedding"] = dense_vec
+
+            # 稀疏
+            if return_sparse:
+                lw = output["lexical_weights"][i]
+                # lexical_weights 是 defaultdict(token_id_str → weight)
+                # 转为 {indices: [int, ...], values: [float, ...]}
+                indices = sorted(int(k) for k in lw.keys())
+                values = [float(lw[str(idx)]) for idx in indices]
+                item["sparse"] = {"indices": indices, "values": values}
+
+            # ColBERT 多向量
+            if return_colbert:
+                colbert_data = output["colbert_vecs"][i]
+                # numpy/torch 矩阵 → list of list
+                if hasattr(colbert_data, "tolist"):
+                    item["colbert"] = colbert_data.tolist()
+                else:
+                    item["colbert"] = [v.tolist() if hasattr(v, "tolist") else list(v) for v in colbert_data]
+
+            results.append(item)
+
+        return results
+
+    async def unload(self) -> None:
+        self._model = None
+        self.model_loaded = False
+        logger.info("Embedding 模型已卸载")
+
+    async def close(self) -> None:
+        await self.unload()

yuanqi/src/engine/llm.py

@@ -0,0 +1,125 @@
+"""
+本地 LLM 推理引擎。
+
+基于 llama-cpp-python 直接加载 GGUF 模型，提供 OpenAI 兼容的 Chat Completions 接口。
+多模态（视觉理解）通过 llava-llama.cpp 的 mmproj 支持。
+
+模型文件统一放在项目根目录 models/gguf/ 下。
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from src.config import settings
+from src.utils.device import detect_device
+from src.utils.logging import logger
+
+
+class LlamaEngine:
+    """本地 GGUF LLM 推理引擎。"""
+
+    name: str = "llm"
+    model_loaded: bool = False
+
+    def __init__(self) -> None:
+        self._model_path = settings.llm_model_path_resolved
+        self._vision_path = settings.vision_model_path_resolved
+        self._device = detect_device(settings.device)
+        self._ctx_size = settings.llm_ctx_size
+        self._llm = None  # llama_cpp.Llama 实例
+        logger.info(
+            "LLM 引擎初始化 | model={} vision={} device={} ctx={}",
+            self._model_path or "(未配置)",
+            self._vision_path or "(未配置)",
+            self._device,
+            self._ctx_size,
+        )
+
+    async def health(self) -> dict[str, Any]:
+        if not self.model_loaded:
+            return {"status": "unavailable", "model": self._model_path or ""}
+        return {"status": "ok", "model": self._model_path}
+
+    async def load(self) -> None:
+        """加载 GGUF 模型（同步操作；llama-cpp-python 的初始化在构造时完成）。"""
+        if self.model_loaded or not self._model_path:
+            return
+        try:
+            import llama_cpp  # noqa: F811
+
+            logger.info("正在加载 LLM 模型: {}", self._model_path)
+            self._llm = llama_cpp.Llama(
+                model_path=self._model_path,
+                n_ctx=self._ctx_size,
+                n_gpu_layers=-1 if self._device in ("cuda", "mps") else 0,
+                verbose=False,
+            )
+            self.model_loaded = True
+            logger.info("LLM 模型加载完成")
+        except ImportError:
+            logger.warning("llama-cpp-python 未安装，LLM 引擎不可用（pip install 'hunjing-yuanqi[llm]'）")
+        except Exception as e:
+            logger.error("LLM 模型加载失败: {}", e)
+
+    async def chat_completion(self, body: dict[str, Any]) -> dict[str, Any]:
+        """执行 Chat Completions 推理。
+
+        Args:
+            body: OpenAI Chat Completions 格式的请求体。
+
+        Returns:
+            OpenAI Chat Completions 格式的响应。
+        """
+        if not self.model_loaded or self._llm is None:
+            return {
+                "error": {
+                    "code": "model_not_loaded",
+                    "message": "LLM 模型未加载，请检查模型路径和依赖安装",
+                }
+            }
+
+        messages = body.get("messages", [])
+        max_tokens = body.get("max_tokens", settings.llm_max_tokens)
+        temperature = body.get("temperature", settings.llm_temperature)
+        stream = body.get("stream", False)
+
+        try:
+            # llama-cpp-python 已内置 OpenAI 兼容格式转换
+            response = self._llm.create_chat_completion(
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stream=stream,
+            )
+
+            if stream:
+                # 对于流式响应，返回 generator
+                return self._stream_response(response)
+            return response
+
+        except Exception as e:
+            logger.error("LLM 推理失败: {}", e)
+            return {
+                "error": {
+                    "code": "inference_error",
+                    "message": str(e),
+                }
+            }
+
+    async def _stream_response(self, response) -> dict[str, Any]:
+        """将 llama-cpp-python 的流式 generator 转为可序列化结构。
+
+        实际使用时（StreamingResponse），在端点层逐块 yield。
+        此方法作为标记返回，端点层特殊处理。
+        """
+        return {"_stream": response}
+
+    async def unload(self) -> None:
+        if self._llm is not None:
+            self._llm = None
+            self.model_loaded = False
+            logger.info("LLM 模型已卸载")
+
+    async def close(self) -> None:
+        await self.unload()

yuanqi/src/engine/reranker.py

@@ -0,0 +1,181 @@
+"""
+本地 Reranker 引擎。
+
+基于 HuggingFace transformers 加载交叉编码器（CrossEncoder）模型
+（如 BAAI/bge-reranker-v2-m3），对查询和文档对进行相关性评分和排序。
+
+使用 AutoModelForSequenceClassification + AutoTokenizer 从 transformers 直接加载。
+
+模型下载位置（按优先级）：
+1. ``YUANQI_RERANKER_MODEL`` 为本地路径时，直接从该路径加载
+2. ``YUANQI_RERANKER_MODEL`` 为 HuggingFace 模型名时，下载到
+   ``models/transformers/<模型名>`` 后再加载
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from src.utils.logging import logger
+import torch
+
+from src.config import settings
+from src.utils.device import detect_device
+
+
+class RerankerEngine:
+    """本地 Reranker 引擎（transformers CrossEncoder）。"""
+
+    name: str = "reranker"
+    model_loaded: bool = False
+
+    def __init__(self) -> None:
+        from src.utils.device import detect_device
+
+        self._device = detect_device(settings.device)
+        self._model_name = settings.reranker_model or "transformers/bge-reranker-v2-m3"
+        self._model = None
+        self._tokenizer = None
+        logger.info(
+            "Reranker 引擎配置 | model={} device={}",
+            self._model_name,
+            self._device,
+        )
+
+    def _resolve_model_path(self) -> str:
+        """解析模型路径。
+
+        若 ``reranker_model`` 是 HuggingFace 模型名时，下载到
+        ``models/transformers/<模型名>`` 后再加载。
+        """
+        # 1) 直接是本地路径
+        direct = Path(self._model_name)
+        if direct.is_absolute() and direct.exists():
+            return str(direct)
+        candidate = settings.models_dir_resolved / self._model_name
+        if candidate.exists():
+            return str(candidate)
+
+        # 2) 已下载到目标目录
+        target_dir = settings.models_dir_resolved / "transformers" / self._model_name
+        if target_dir.exists():
+            return str(target_dir)
+
+        # 3) 已下载到 models/transformers/<repo_name>
+        repo_name = Path(self._model_name).name
+        target_dir2 = settings.models_dir_resolved / "transformers" / repo_name
+        if target_dir2.exists():
+            return str(target_dir2)
+
+        # 4) 下载到 models/transformers/<repo_name>
+        target = target_dir2
+        logger.info("本地模型不存在，从 {} 下载到 {}", settings.model_source, target)
+        from src.utils.downloader import download_model
+
+        if download_model(self._model_name, target, settings.model_source):
+            return str(target)
+        logger.warning("模型下载失败，回退到 HF 缓存加载")
+        return self._model_name
+
+    async def health(self) -> dict[str, Any]:
+        if not self.model_loaded:
+            return {"status": "unavailable", "model": self._model_name}
+        return {"status": "ok", "model": self._model_name}
+
+    async def load(self) -> None:
+        """加载 transformers 交叉编码器模型。
+
+        优先使用本地路径；本地不存在时用 HuggingFace 模型名首次下载到 HF 缓存。
+        """
+        if self.model_loaded:
+            return
+        try:
+            model_path = self._resolve_model_path()
+            from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+            logger.info("正在加载 Reranker 模型: {}", model_path)
+            self._tokenizer = AutoTokenizer.from_pretrained(model_path)
+            self._model = AutoModelForSequenceClassification.from_pretrained(model_path)
+            self._model.eval()
+            self._model.to(self._device)
+            self.model_loaded = True
+            logger.info("Reranker 模型加载完成: {}", model_path)
+        except Exception as e:
+            logger.error("Reranker 模型加载失败: {}", e)
+
+    def count_tokens(self, query: str, documents: list[str]) -> dict:
+        """统计 reranker 完整 token 用量。
+
+        Returns:
+            含 ``input``（query_tokens / document_tokens[] / total_tokens）
+            和 ``compute``（pairs / tokens）的嵌套 dict。
+        """
+        from src.utils.token_counter import RerankerTokenCountStrategy
+
+        if self._tokenizer is None:
+            return {
+                "input": {
+                    "query_tokens": 0,
+                    "document_tokens": [0] * len(documents),
+                    "total_tokens": 0,
+                },
+                "compute": {"pairs": len(documents), "tokens": 0},
+            }
+        strategy = RerankerTokenCountStrategy()
+        usage = strategy.count(self._tokenizer, query=query, documents=documents)
+        return usage.to_dict()
+
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        top_n: int | None = None,
+    ) -> list[dict[str, Any]]:
+        """对文档执行重排序。
+
+        Args:
+            query: 搜索查询文本。
+            documents: 待排序文档列表。
+            top_n: 返回前 N 个结果（默认返回全部排序）。
+
+        Returns:
+            按相关性降序排列的 [{index, document, relevance_score}, ...] 列表。
+        """
+        if not self.model_loaded or self._model is None or self._tokenizer is None:
+            return [
+                {"index": i, "document": doc, "relevance_score": 0.0}
+                for i, doc in enumerate(documents)
+            ]
+
+        pairs = [[query, doc] for doc in documents]
+        inputs = self._tokenizer(
+            pairs,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+            max_length=512,
+        ).to(self._device)
+
+        with torch.no_grad():
+            outputs = self._model(**inputs)
+            scores = outputs.logits.sigmoid().flatten().tolist()
+
+        scored = [
+            {"index": i, "document": documents[i], "relevance_score": float(score)}
+            for i, score in enumerate(scores)
+        ]
+        scored.sort(key=lambda x: x["relevance_score"], reverse=True)
+
+        if top_n is not None and top_n < len(scored):
+            scored = scored[:top_n]
+        return scored
+
+    async def unload(self) -> None:
+        self._model = None
+        self._tokenizer = None
+        self.model_loaded = False
+        logger.info("Reranker 模型已卸载")
+
+    async def close(self) -> None:
+        await self.unload()