anti-plagiarism/services/worker-gpu/app/faiss_manager.py

"""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