7758315632
Services: - api: FastAPI gateway with JWT auth, async endpoints, WebSocket - worker-gpu: CUDA sentence-transformers, FAISS IVFFlat, Ollama LLM - worker-indexer: Winnowing+MinHash plagiarism detection, PDF/DOCX extraction - worker-notifier: SMTP email notifications - worker-gost: GOST 7.1-2003 and GOST R 7.0.5-2008 formatting Infrastructure: - docker-compose.yml (production) + docker-compose.dev.yml (hot reload) - Nginx reverse proxy + WebSocket support - PostgreSQL 16 with Alembic migrations - Elasticsearch 8 with Russian/English analyzers - MinIO, RabbitMQ, Redis, Ollama Frontend: - React 18 + Vite + TypeScript + TailwindCSS + Zustand + React Query v5 - 9 pages: Home, Search, Cabinet, Task, Check, Bibliography, Pricing, Login, Register Scripts: - Parser stubs: OpenAlex, КиберЛенинка, arXiv (Phase 0 - to be filled) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
82 lines
3.0 KiB
Python
82 lines
3.0 KiB
Python
"""Singleton менеджер sentence-transformers модели.
|
||
|
||
Модель загружается один раз при первом обращении и кэшируется в памяти GPU.
|
||
"""
|
||
|
||
import logging
|
||
|
||
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), 768), нормализованный для cosine similarity
|
||
"""
|
||
if not texts:
|
||
return np.array([]).reshape(0, settings.EMBED_DIM)
|
||
|
||
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_single(cls, text: str) -> np.ndarray:
|
||
"""Закодировать один текст. Удобный метод."""
|
||
return cls.encode([text])[0]
|
||
|
||
@classmethod
|
||
def unload(cls) -> None:
|
||
"""Выгрузить модель из памяти (для тестов/диагностики)."""
|
||
cls._model = None
|
||
logger.info("Модель выгружена из памяти")
|