232 lines
7.6 KiB
Python
232 lines
7.6 KiB
Python
"""嵌入模型抽象层 — 策略模式.
|
||
|
||
支持的 Provider:
|
||
openai — OpenAI / 硅基流动 / 智谱 / DeepSeek / 月之暗面 等 OpenAI 兼容服务
|
||
dashscope — 阿里云 DashScope (通义千问)
|
||
|
||
使用方式:
|
||
from src.core.config import EmbedConfig
|
||
from src.core.embedder import create_embedder, batch_embed
|
||
|
||
# 本地模型
|
||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||
# OpenAI 兼容 API
|
||
embedder = create_embedder(EmbedConfig(mode="api", provider="openai"))
|
||
# 阿里云 DashScope
|
||
embedder = create_embedder(EmbedConfig(mode="api", provider="dashscope"))
|
||
|
||
vectors = batch_embed(embedder, long_text_list)
|
||
"""
|
||
import logging
|
||
import os
|
||
import threading
|
||
from typing import Protocol
|
||
|
||
from src.core.config import EmbedConfig
|
||
|
||
logger = logging.getLogger("md-vector-db")
|
||
|
||
_HF_MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com")
|
||
|
||
_HF_ENV_LOCK = threading.Lock()
|
||
|
||
# -- Provider 默认配置 --
|
||
_PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = {
|
||
"openai": {
|
||
"api_base": "https://api.openai.com/v1",
|
||
"model": "text-embedding-3-small",
|
||
"dimension": 1536,
|
||
},
|
||
"dashscope": {
|
||
"api_base": "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding",
|
||
"model": "text-embedding-v4",
|
||
"dimension": 1536,
|
||
},
|
||
}
|
||
|
||
|
||
# -- 接口 --
|
||
class Embedder(Protocol):
|
||
"""嵌入器接口."""
|
||
|
||
@property
|
||
def dimension(self) -> int:
|
||
"""返回嵌入向量的维度."""
|
||
...
|
||
|
||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||
"""对文本列表进行嵌入.
|
||
|
||
Args:
|
||
texts: 待嵌入的文本列表
|
||
|
||
Returns:
|
||
嵌入向量列表，每个向量为 float 列表
|
||
"""
|
||
...
|
||
|
||
|
||
# -- 本地模型 --
|
||
class LocalEmbedder:
|
||
"""sentence-transformers 本地模型嵌入器. 自动检测 GPU."""
|
||
|
||
def __init__(self, config: EmbedConfig):
|
||
from sentence_transformers import SentenceTransformer
|
||
|
||
# 自动检测 GPU
|
||
try:
|
||
import torch
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
if device == "cuda":
|
||
logger.info("检测到 GPU: %s", torch.cuda.get_device_name(0))
|
||
except ImportError:
|
||
device = "cpu"
|
||
|
||
self._config = config
|
||
try:
|
||
self._model = SentenceTransformer(
|
||
config.local_model, local_files_only=True, device=device
|
||
)
|
||
except Exception:
|
||
logger.info("模型未缓存, 通过镜像下载 %s", config.local_model)
|
||
# 通过 HF_ENDPOINT 环境变量设置镜像（sentence-transformers 依赖 huggingface_hub）
|
||
# 临时设置仅用于模型下载，下载完成后还原
|
||
old_endpoint = os.environ.get("HF_ENDPOINT")
|
||
with _HF_ENV_LOCK:
|
||
os.environ["HF_ENDPOINT"] = _HF_MIRROR
|
||
try:
|
||
self._model = SentenceTransformer(
|
||
config.local_model, device=device
|
||
)
|
||
finally:
|
||
with _HF_ENV_LOCK:
|
||
if old_endpoint is not None:
|
||
os.environ["HF_ENDPOINT"] = old_endpoint
|
||
else:
|
||
os.environ.pop("HF_ENDPOINT", None)
|
||
|
||
@property
|
||
def dimension(self) -> int:
|
||
return self._model.get_embedding_dimension()
|
||
|
||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||
if not texts:
|
||
raise ValueError("文本列表不能为空")
|
||
embeddings = self._model.encode(texts, normalize_embeddings=True)
|
||
return embeddings.tolist()
|
||
|
||
|
||
# -- API Provider 基类 --
|
||
class _BaseAPIEmbedder:
|
||
"""API 嵌入器基类: 统一 api_base/model 解析逻辑."""
|
||
|
||
def __init__(self, config: EmbedConfig, provider: str):
|
||
defaults = _PROVIDER_DEFAULTS.get(provider, {})
|
||
self._api_base = config.api_base or str(defaults.get("api_base", ""))
|
||
self._model = config.model or str(defaults.get("model", ""))
|
||
self._api_key = config.api_key
|
||
self._dimension = int(defaults.get("dimension", 1536))
|
||
|
||
@property
|
||
def dimension(self) -> int:
|
||
return self._dimension
|
||
|
||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||
raise NotImplementedError
|
||
|
||
|
||
# -- OpenAI 兼容 API --
|
||
class OpenAIEmbedder(_BaseAPIEmbedder):
|
||
"""OpenAI / 硅基流动 / 智谱 / DeepSeek / 月之暗面 等."""
|
||
|
||
def __init__(self, config: EmbedConfig):
|
||
super().__init__(config, "openai")
|
||
self._client = None # 懒初始化，首次 embed() 时创建
|
||
|
||
def _ensure_client(self):
|
||
"""懒创建 OpenAI 客户端（避免 import 时依赖 openai 包）."""
|
||
if self._client is None:
|
||
from openai import OpenAI
|
||
self._client = OpenAI(base_url=self._api_base, api_key=self._api_key)
|
||
|
||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||
if not texts:
|
||
raise ValueError("文本列表不能为空")
|
||
self._ensure_client()
|
||
response = self._client.embeddings.create(model=self._model, input=texts)
|
||
return [d.embedding for d in response.data]
|
||
|
||
|
||
# -- 阿里云 DashScope --
|
||
class DashscopeEmbedder(_BaseAPIEmbedder):
|
||
"""阿里云 DashScope 嵌入 (自定义 HTTP API, 非 OpenAI 兼容)."""
|
||
|
||
def __init__(self, config: EmbedConfig):
|
||
super().__init__(config, "dashscope")
|
||
|
||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||
if not texts:
|
||
raise ValueError("文本列表不能为空")
|
||
import requests # 惰性导入（仅 DashScope 使用）
|
||
resp = requests.post(
|
||
self._api_base,
|
||
headers={
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json={
|
||
"model": self._model,
|
||
"input": {"texts": texts},
|
||
},
|
||
timeout=60,
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
# DashScope 返回: {"output": {"embeddings": [{"text_index": 0, "embedding": [...]}, ...]}}
|
||
output = data.get("output")
|
||
if output is None:
|
||
raise ValueError(f"DashScope 响应缺少 output 字段: {data}")
|
||
embeddings_raw = output.get("embeddings")
|
||
if not isinstance(embeddings_raw, list):
|
||
raise ValueError(f"DashScope embeddings 不是列表: {type(embeddings_raw)}")
|
||
# 按 text_index 排序确保顺序
|
||
embeddings_raw_sorted = sorted(embeddings_raw, key=lambda x: x.get("text_index", 0))
|
||
return [e["embedding"] for e in embeddings_raw_sorted]
|
||
|
||
|
||
# -- 工厂函数 --
|
||
_PROVIDER_CLASSES: dict[str, type[_BaseAPIEmbedder]] = {
|
||
"openai": OpenAIEmbedder,
|
||
"dashscope": DashscopeEmbedder,
|
||
}
|
||
|
||
SUPPORTED_PROVIDERS = list(_PROVIDER_CLASSES.keys())
|
||
|
||
|
||
def create_embedder(config: EmbedConfig) -> Embedder:
|
||
"""工厂函数: 根据配置创建嵌入器."""
|
||
if config.mode == "local":
|
||
return LocalEmbedder(config)
|
||
|
||
if config.mode == "api":
|
||
provider = config.provider or "openai"
|
||
cls = _PROVIDER_CLASSES.get(provider)
|
||
if cls is None:
|
||
raise ValueError(
|
||
f"不支持的 provider: {provider}, 可选: {SUPPORTED_PROVIDERS}"
|
||
)
|
||
return cls(config)
|
||
|
||
raise ValueError(f"不支持的嵌入模式: {config.mode}")
|
||
|
||
|
||
def batch_embed(
|
||
embedder: Embedder, texts: list[str], batch_size: int = 32
|
||
) -> list[list[float]]:
|
||
"""分批嵌入, 避免一次性传入过多文本导致 OOM."""
|
||
all_embeddings = []
|
||
for i in range(0, len(texts), batch_size):
|
||
batch = texts[i:i + batch_size]
|
||
all_embeddings.extend(embedder.embed(batch))
|
||
return all_embeddings
|