feat: initial microservices project structure
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>
This commit is contained in:
jze9
210
services/worker-gpu/app/faiss_manager.py
Normal file
210
services/worker-gpu/app/faiss_manager.py
Normal file
@@ -0,0 +1,210 @@
|
||||
"""Singleton менеджер FAISS GPU индекса.
|
||||
|
||||
Использует IVFFlat (а не HNSW — не поддерживается на GPU).
|
||||
Поддерживает graceful degradation на CPU если GPU недоступна.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FAISSManager:
|
||||
"""Singleton для управления FAISS индексом на GPU."""
|
||||
|
||||
_index = None
|
||||
_id_map: dict[int, int] = {} # faiss_internal_id -> doc_id (PostgreSQL)
|
||||
_reverse_map: dict[int, int] = {} # doc_id -> faiss_internal_id
|
||||
_use_gpu: bool = False
|
||||
_is_trained: bool = False
|
||||
|
||||
@classmethod
|
||||
def load_or_create(cls) -> None:
|
||||
"""
|
||||
Загрузить индекс с диска или создать новый.
|
||||
|
||||
Пытается перенести индекс на GPU, при ошибке остаётся на CPU.
|
||||
"""
|
||||
import faiss
|
||||
|
||||
index_path = settings.FAISS_INDEX_PATH
|
||||
id_map_path = settings.FAISS_ID_MAP_PATH
|
||||
|
||||
# Загрузить ID маппинг
|
||||
if os.path.exists(id_map_path):
|
||||
with open(id_map_path, "r") as f:
|
||||
raw_map = json.load(f)
|
||||
cls._id_map = {int(k): int(v) for k, v in raw_map.items()}
|
||||
cls._reverse_map = {v: k for k, v in cls._id_map.items()}
|
||||
logger.info(f"ID маппинг загружен: {len(cls._id_map)} записей")
|
||||
|
||||
if os.path.exists(index_path):
|
||||
# Загрузить существующий индекс
|
||||
logger.info(f"Загрузка FAISS индекса из {index_path}")
|
||||
cpu_index = faiss.read_index(index_path)
|
||||
cls._is_trained = cpu_index.is_trained
|
||||
else:
|
||||
# Создать новый IVFFlat индекс
|
||||
logger.info("Создание нового FAISS IVFFlat индекса...")
|
||||
quantizer = faiss.IndexFlatIP(settings.EMBED_DIM)
|
||||
cpu_index = faiss.IndexIVFFlat(
|
||||
quantizer,
|
||||
settings.EMBED_DIM,
|
||||
settings.FAISS_NLIST,
|
||||
faiss.METRIC_INNER_PRODUCT,
|
||||
)
|
||||
# IVFFlat требует обучения перед использованием
|
||||
cls._is_trained = False
|
||||
|
||||
# Попытка перенести на GPU
|
||||
try:
|
||||
res = faiss.StandardGpuResources()
|
||||
cls._index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
|
||||
cls._use_gpu = True
|
||||
logger.info("FAISS индекс размещён на GPU")
|
||||
except Exception as e:
|
||||
logger.warning(f"GPU недоступна: {e}. Используем CPU FAISS.")
|
||||
cls._index = cpu_index
|
||||
cls._use_gpu = False
|
||||
|
||||
if cls._is_trained:
|
||||
cls._index.nprobe = settings.FAISS_NPROBE
|
||||
|
||||
@classmethod
|
||||
def search(cls, query_vector: np.ndarray, k: int = 20) -> list[tuple[int, float]]:
|
||||
"""
|
||||
Поиск k ближайших векторов.
|
||||
|
||||
Args:
|
||||
query_vector: Нормализованный вектор запроса, форма (768,)
|
||||
k: Количество результатов
|
||||
|
||||
Returns:
|
||||
Список кортежей (doc_id, cosine_score), отсортированных по убыванию score
|
||||
"""
|
||||
if cls._index is None or not cls._is_trained:
|
||||
logger.warning("FAISS индекс не инициализирован или не обучен, пропускаем поиск")
|
||||
return []
|
||||
|
||||
try:
|
||||
query = query_vector.reshape(1, -1).astype(np.float32)
|
||||
distances, indices = cls._index.search(query, k)
|
||||
|
||||
results = []
|
||||
for idx, dist in zip(indices[0], distances[0]):
|
||||
if idx == -1:
|
||||
continue
|
||||
doc_id = cls._id_map.get(int(idx))
|
||||
if doc_id is not None:
|
||||
results.append((doc_id, float(dist)))
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка поиска FAISS: {e}")
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def add_vectors(cls, vectors: np.ndarray, doc_ids: list[int]) -> None:
|
||||
"""
|
||||
Добавить векторы в индекс.
|
||||
|
||||
Если индекс не обучен и накопилось достаточно векторов — обучить его.
|
||||
|
||||
Args:
|
||||
vectors: numpy массив формы (N, 768)
|
||||
doc_ids: Список doc_id из PostgreSQL
|
||||
"""
|
||||
import faiss
|
||||
|
||||
if cls._index is None:
|
||||
cls.load_or_create()
|
||||
|
||||
vectors = vectors.astype(np.float32)
|
||||
|
||||
if not cls._is_trained:
|
||||
# Для IVFFlat нужно минимум nlist * 39 обучающих примеров
|
||||
min_train = settings.FAISS_NLIST * 39
|
||||
current_n = cls._index.ntotal if cls._index is not None else 0
|
||||
|
||||
if current_n + len(vectors) >= min_train:
|
||||
logger.info(f"Обучение IVFFlat индекса на {current_n + len(vectors)} векторах...")
|
||||
cls._index.train(vectors)
|
||||
cls._is_trained = True
|
||||
cls._index.nprobe = settings.FAISS_NPROBE
|
||||
logger.info("Обучение завершено")
|
||||
else:
|
||||
logger.info(
|
||||
f"Недостаточно векторов для обучения IVFFlat "
|
||||
f"({current_n + len(vectors)} < {min_train}). "
|
||||
"Используйте FlatIP до накопления достаточного количества документов."
|
||||
)
|
||||
# Временный flat индекс для малого количества документов
|
||||
if not hasattr(cls, '_flat_index') or cls._flat_index is None:
|
||||
cls._flat_index = faiss.IndexFlatIP(settings.EMBED_DIM)
|
||||
|
||||
# Добавить в flat индекс
|
||||
start_id = cls._flat_index.ntotal
|
||||
cls._flat_index.add(vectors)
|
||||
for i, doc_id in enumerate(doc_ids):
|
||||
internal_id = start_id + i
|
||||
cls._id_map[internal_id] = doc_id
|
||||
cls._reverse_map[doc_id] = internal_id
|
||||
cls._save_id_map()
|
||||
return
|
||||
|
||||
if cls._is_trained:
|
||||
start_id = cls._index.ntotal
|
||||
cls._index.add(vectors)
|
||||
for i, doc_id in enumerate(doc_ids):
|
||||
internal_id = start_id + i
|
||||
cls._id_map[internal_id] = doc_id
|
||||
cls._reverse_map[doc_id] = internal_id
|
||||
|
||||
cls._save_id_map()
|
||||
logger.info(f"Добавлено {len(doc_ids)} векторов в FAISS. Всего: {cls._index.ntotal}")
|
||||
|
||||
@classmethod
|
||||
def save(cls) -> None:
|
||||
"""Сохранить индекс на диск (CPU версия)."""
|
||||
if cls._index is None:
|
||||
return
|
||||
|
||||
import faiss
|
||||
|
||||
index_path = Path(settings.FAISS_INDEX_PATH)
|
||||
index_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Перенести на CPU перед сохранением
|
||||
if cls._use_gpu:
|
||||
cpu_index = faiss.index_gpu_to_cpu(cls._index)
|
||||
else:
|
||||
cpu_index = cls._index
|
||||
|
||||
faiss.write_index(cpu_index, str(index_path))
|
||||
cls._save_id_map()
|
||||
logger.info(f"FAISS индекс сохранён: {index_path} ({cpu_index.ntotal} векторов)")
|
||||
|
||||
@classmethod
|
||||
def _save_id_map(cls) -> None:
|
||||
"""Сохранить маппинг faiss_id -> doc_id на диск."""
|
||||
id_map_path = Path(settings.FAISS_ID_MAP_PATH)
|
||||
id_map_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(id_map_path, "w") as f:
|
||||
json.dump({str(k): v for k, v in cls._id_map.items()}, f)
|
||||
|
||||
@classmethod
|
||||
@property
|
||||
def total_vectors(cls) -> int:
|
||||
"""Количество векторов в индексе."""
|
||||
if cls._index is None:
|
||||
return 0
|
||||
return cls._index.ntotal
|
||||
Reference in New Issue
Block a user