c156fb2b78
EMBED_BACKEND=cloud — эмбеддинги через DeepInfra (OpenAI-совместимый API), модель та же BAAI/bge-m3, что уже льётся на VM109 — переключение не рвёт совместимость с уже пересчитанным FAISS-индексом (тот же вектор 1024). Не активно по умолчанию (EMBED_BACKEND всё ещё "ollama").
138 lines
5.8 KiB
Python
138 lines
5.8 KiB
Python
"""Singleton менеджер эмбеддинг-модели.
|
||
|
||
Бэкенд переключается через EMBED_BACKEND:
|
||
- "ollama" — HTTP-вызов к Ollama /api/embed (GGUF-модель через Vulkan, работает на AMD GPU)
|
||
- "sentence_transformers" — локальная загрузка через sentence-transformers (CUDA/CPU)
|
||
- "cloud" — облачный инференс той же модели (DeepInfra, OpenAI-совместимый формат)
|
||
|
||
Модель (в случае sentence_transformers) загружается один раз при первом обращении
|
||
и кэшируется в памяти GPU.
|
||
"""
|
||
|
||
import logging
|
||
|
||
import httpx
|
||
import numpy as np
|
||
|
||
from app.config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ModelManager:
|
||
"""Singleton для управления эмбеддинг-моделью."""
|
||
|
||
_instance = None
|
||
_model = None
|
||
_device: str = settings.EMBED_DEVICE
|
||
|
||
@classmethod
|
||
def get_model(cls):
|
||
"""Получить или загрузить модель sentence-transformers."""
|
||
if cls._model is None:
|
||
# Импорт здесь, чтобы не блокировать импорт модуля если torch не установлен
|
||
from sentence_transformers import SentenceTransformer
|
||
|
||
logger.info(
|
||
f"Загрузка модели {settings.EMBED_MODEL!r} на устройство {settings.EMBED_DEVICE!r}..."
|
||
)
|
||
try:
|
||
cls._model = SentenceTransformer(
|
||
settings.EMBED_MODEL,
|
||
device=settings.EMBED_DEVICE,
|
||
)
|
||
logger.info(
|
||
f"Модель загружена. Размерность вектора: {cls._model.get_sentence_embedding_dimension()}"
|
||
)
|
||
except Exception as e:
|
||
# Graceful degradation на CPU если CUDA недоступна
|
||
logger.warning(f"Не удалось загрузить модель на CUDA: {e}. Переключение на CPU.")
|
||
cls._device = "cpu"
|
||
cls._model = SentenceTransformer(settings.EMBED_MODEL, device="cpu")
|
||
|
||
return cls._model
|
||
|
||
@classmethod
|
||
def encode(cls, texts: list[str]) -> np.ndarray:
|
||
"""
|
||
Закодировать тексты в векторы.
|
||
|
||
Args:
|
||
texts: Список текстов для кодирования
|
||
|
||
Returns:
|
||
numpy массив формы (len(texts), EMBED_DIM), нормализованный для cosine similarity
|
||
"""
|
||
if not texts:
|
||
return np.array([]).reshape(0, settings.EMBED_DIM)
|
||
|
||
if settings.EMBED_BACKEND == "ollama":
|
||
return cls._encode_ollama(texts)
|
||
if settings.EMBED_BACKEND == "cloud":
|
||
return cls._encode_cloud(texts)
|
||
|
||
model = cls.get_model()
|
||
vectors = model.encode(
|
||
texts,
|
||
batch_size=settings.EMBED_BATCH_SIZE,
|
||
normalize_embeddings=True, # Нормализация для cosine через inner product
|
||
show_progress_bar=len(texts) > 100,
|
||
convert_to_numpy=True,
|
||
)
|
||
return vectors.astype(np.float32)
|
||
|
||
@classmethod
|
||
def _encode_ollama(cls, texts: list[str]) -> np.ndarray:
|
||
"""Закодировать тексты через Ollama /api/embed (пакетами по EMBED_BATCH_SIZE)."""
|
||
all_vectors: list[list[float]] = []
|
||
batch_size = settings.EMBED_BATCH_SIZE
|
||
for i in range(0, len(texts), batch_size):
|
||
batch = texts[i : i + batch_size]
|
||
response = httpx.post(
|
||
f"{settings.OLLAMA_URL}/api/embed",
|
||
json={"model": settings.EMBED_MODEL, "input": batch},
|
||
timeout=120.0,
|
||
)
|
||
response.raise_for_status()
|
||
all_vectors.extend(response.json()["embeddings"])
|
||
|
||
vectors = np.array(all_vectors, dtype=np.float32)
|
||
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||
norms[norms == 0] = 1.0
|
||
return vectors / norms # Нормализация для cosine через inner product
|
||
|
||
@classmethod
|
||
def _encode_cloud(cls, texts: list[str]) -> np.ndarray:
|
||
"""Закодировать тексты через облачный инференс (DeepInfra, OpenAI-формат)."""
|
||
all_vectors: list[list[float]] = []
|
||
batch_size = settings.EMBED_BATCH_SIZE
|
||
for i in range(0, len(texts), batch_size):
|
||
batch = texts[i : i + batch_size]
|
||
response = httpx.post(
|
||
settings.CLOUD_EMBED_URL,
|
||
json={"model": settings.CLOUD_EMBED_MODEL, "input": batch},
|
||
headers={"Authorization": f"Bearer {settings.CLOUD_EMBED_API_KEY}"},
|
||
timeout=120.0,
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()["data"]
|
||
# OpenAI-формат может вернуть элементы не по порядку — сортируем по index
|
||
data.sort(key=lambda d: d["index"])
|
||
all_vectors.extend(d["embedding"] for d in data)
|
||
|
||
vectors = np.array(all_vectors, dtype=np.float32)
|
||
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||
norms[norms == 0] = 1.0
|
||
return vectors / norms # Нормализация для cosine через inner product
|
||
|
||
@classmethod
|
||
def encode_single(cls, text: str) -> np.ndarray:
|
||
"""Закодировать один текст. Удобный метод."""
|
||
return cls.encode([text])[0]
|
||
|
||
@classmethod
|
||
def unload(cls) -> None:
|
||
"""Выгрузить модель из памяти (для тестов/диагностики)."""
|
||
cls._model = None
|
||
logger.info("Модель выгружена из памяти")
|