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