Video-Tree-TRM5/adapters/embedding.py

"""嵌入适配器 —— local/remote 双后端实现。

封装文本嵌入器，支持本地 sentence-transformers 和远程 OpenAI 兼容 API 两种后端。
提供统一的 ``embed()`` / ``embed_tensor()`` 接口，冻结不训练。
两个类均满足 ``app.ports.EmbeddingProvider`` Protocol。
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import torch
from loguru import logger

if TYPE_CHECKING:
    from numpy import ndarray
    from torch import Tensor


class LocalEmbeddingProvider:
    """本地 sentence-transformers 嵌入器（冻结）。

    使用 HuggingFace sentence-transformers 加载模型进行本地推理，
    所有参数冻结，仅用于嵌入提取。

    属性:
        dim: 嵌入维度 D。
    """

    def __init__(self, model_name: str, embed_dim: int, device: str = "cpu") -> None:
        """初始化本地嵌入模型。

        参数:
            model_name: HuggingFace 模型名称（如 'BAAI/bge-base-zh-v1.5'）。
            embed_dim: 期望的嵌入维度。
            device: 推理设备（'cpu' / 'cuda' 等）。

        异常:
            AssertionError: 模型实际维度与 embed_dim 不一致。
        """
        from sentence_transformers import SentenceTransformer

        self._dim = embed_dim

        self._model = SentenceTransformer(model_name, device=device)
        self._model.eval()
        # 冻结所有参数
        for param in self._model.parameters():
            param.requires_grad = False

        actual_dim = self._model.get_sentence_embedding_dimension()
        assert actual_dim == self._dim, (
            f"模型实际维度 ({actual_dim}) 与配置 embed_dim ({self._dim}) 不一致"
        )

        logger.info("本地嵌入模型初始化完成", model=model_name, device=device)

    # ------------------------------------------------------------------
    # 公共接口
    # ------------------------------------------------------------------

    @property
    def dim(self) -> int:
        """嵌入维度 D。"""
        return self._dim

    def embed(self, texts: str | list[str]) -> ndarray:
        """文本 → 嵌入向量（L2 归一化）。

        参数:
            texts: 单条文本或文本列表。

        返回:
            [N, D] ndarray，每行 L2 范数为 1.0。单条文本时 N=1。
        """
        if isinstance(texts, str):
            texts = [texts]

        with torch.no_grad():
            embeddings = self._model.encode(
                texts,
                normalize_embeddings=True,
                convert_to_numpy=True,
            )
        # sentence-transformers encode 返回 ndarray [N, D]
        if embeddings.ndim == 1:
            embeddings = embeddings.reshape(1, -1)
        return embeddings

    def embed_tensor(self, texts: str | list[str]) -> Tensor:
        """文本 → 嵌入 Tensor（L2 归一化）。

        参数:
            texts: 单条文本或文本列表。

        返回:
            [N, D] torch.Tensor（float32）。
        """
        arr = self.embed(texts)
        return torch.from_numpy(arr).float()


class RemoteEmbeddingProvider:
    """远程 OpenAI 兼容 API 嵌入器。

    通过 OpenAI 兼容 API（如 GPUStack）调用远程嵌入模型。

    属性:
        dim: 嵌入维度 D。
    """

    def __init__(self, model_name: str, embed_dim: int, api_key: str, api_url: str) -> None:
        """初始化远程嵌入客户端。

        参数:
            model_name: 远程模型名称。
            embed_dim: 期望的嵌入维度。
            api_key: API 密钥。
            api_url: API 基础 URL。

        异常:
            ValueError: api_key 或 api_url 为空。
        """
        if not api_key:
            raise ValueError("远程模式必须提供 api_key")
        if not api_url:
            raise ValueError("远程模式必须提供 api_url")

        from openai import OpenAI

        self._dim = embed_dim
        self._model_name = model_name
        self._client = OpenAI(base_url=api_url, api_key=api_key)

        logger.info("远程嵌入客户端初始化完成", model=model_name, api_url=api_url)

    # ------------------------------------------------------------------
    # 公共接口
    # ------------------------------------------------------------------

    @property
    def dim(self) -> int:
        """嵌入维度 D。"""
        return self._dim

    def embed(self, texts: str | list[str]) -> ndarray:
        """文本 → 嵌入向量（L2 归一化）。

        参数:
            texts: 单条文本或文本列表。

        返回:
            [N, D] ndarray，每行 L2 范数为 1.0。单条文本时 N=1。
        """
        if isinstance(texts, str):
            texts = [texts]

        response = self._client.embeddings.create(
            model=self._model_name,
            input=texts,
        )
        # 按 index 排序，确保顺序一致
        sorted_data = sorted(response.data, key=lambda x: x.index)
        embeddings = np.array([item.embedding for item in sorted_data], dtype=np.float32)

        # L2 归一化
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms = np.maximum(norms, 1e-12)  # 避免除零
        embeddings = embeddings / norms

        return embeddings

    def embed_tensor(self, texts: str | list[str]) -> Tensor:
        """文本 → 嵌入 Tensor（L2 归一化）。

        参数:
            texts: 单条文本或文本列表。

        返回:
            [N, D] torch.Tensor（float32）。
        """
        arr = self.embed(texts)
        return torch.from_numpy(arr).float()