feat(adapters): EmbeddingProvider Protocol + local/remote 双后端实现
- app/ports.py: 添加 EmbeddingProvider Protocol（runtime_checkable，dim 属性 + embed 方法） - adapters/embedding.py: 从参考代码迁移，拆分为 LocalEmbeddingProvider 和 RemoteEmbeddingProvider - Local: sentence-transformers 冻结推理，维度校验 - Remote: OpenAI 兼容 API，L2 归一化，按 index 排序 - 两者均提供 embed() 和 embed_tensor() 统一接口 - tests/unit/test_embedding_adapter.py: Protocol 满足性、形状校验、导入测试 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,184 @@
|
||||
"""嵌入适配器 —— 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()
|
||||
Reference in New Issue
Block a user