91 lines
3.3 KiB
Python
91 lines
3.3 KiB
Python
"""本地 Sentence Transformers 模型加载与推理。"""
|
|
import asyncio
|
|
import json
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List
|
|
|
|
|
|
class LocalEmbeddingService:
|
|
"""从 backend/models 安全加载本地向量模型，并复用已加载实例。"""
|
|
|
|
MODELS_DIR = Path(__file__).resolve().parents[2] / "models"
|
|
_models: Dict[str, Any] = {}
|
|
_load_lock = threading.Lock()
|
|
|
|
@classmethod
|
|
def list_available_models(cls) -> List[Dict[str, Any]]:
|
|
models_root = cls.MODELS_DIR.resolve()
|
|
if not models_root.is_dir():
|
|
return []
|
|
|
|
models = []
|
|
for model_path in sorted(models_root.iterdir(), key=lambda item: item.name.lower()):
|
|
if not model_path.is_dir() or model_path.name.startswith("."):
|
|
continue
|
|
dimension = None
|
|
pooling_config = model_path / "1_Pooling" / "config.json"
|
|
try:
|
|
with pooling_config.open("r", encoding="utf-8") as file:
|
|
dimension = json.load(file).get("word_embedding_dimension")
|
|
except (OSError, ValueError, AttributeError):
|
|
pass
|
|
ready = any(
|
|
(model_path / filename).is_file()
|
|
for filename in ("model.safetensors", "pytorch_model.bin")
|
|
)
|
|
models.append({
|
|
"name": model_path.name,
|
|
"dimension": dimension,
|
|
"ready": ready,
|
|
})
|
|
return models
|
|
|
|
@classmethod
|
|
def resolve_model_path(cls, model_name: str) -> Path:
|
|
name = (model_name or "").strip()
|
|
if not name:
|
|
raise ValueError("缺少本地模型目录名称")
|
|
|
|
models_root = cls.MODELS_DIR.resolve()
|
|
requested = Path(name)
|
|
candidate = requested.resolve() if requested.is_absolute() else (models_root / requested).resolve()
|
|
if not candidate.is_relative_to(models_root):
|
|
raise ValueError("本地模型必须位于 backend/models 目录中")
|
|
if not candidate.is_dir():
|
|
raise ValueError(f"本地模型目录不存在：{candidate.name}")
|
|
return candidate
|
|
|
|
@classmethod
|
|
def _get_model(cls, model_path: Path):
|
|
cache_key = str(model_path)
|
|
with cls._load_lock:
|
|
model = cls._models.get(cache_key)
|
|
if model is not None:
|
|
return model
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
except ImportError as exc:
|
|
raise ValueError(
|
|
"本地向量模型依赖未安装，请执行 pip install -r requirements.txt"
|
|
) from exc
|
|
model = SentenceTransformer(cache_key, local_files_only=True)
|
|
cls._models[cache_key] = model
|
|
return model
|
|
|
|
@classmethod
|
|
def _encode(cls, model_path: Path, text: str) -> List[float]:
|
|
model = cls._get_model(model_path)
|
|
vector = model.encode(
|
|
text,
|
|
normalize_embeddings=True,
|
|
convert_to_numpy=True,
|
|
show_progress_bar=False,
|
|
)
|
|
return [float(value) for value in vector.tolist()]
|
|
|
|
@classmethod
|
|
async def generate_embedding(cls, model_name: str, text: str) -> List[float]:
|
|
model_path = cls.resolve_model_path(model_name)
|
|
return await asyncio.to_thread(cls._encode, model_path, text)
|