diff --git a/services/worker-gpu/app/faiss_manager.py b/services/worker-gpu/app/faiss_manager.py index 1129d3a..90356bd 100644 --- a/services/worker-gpu/app/faiss_manager.py +++ b/services/worker-gpu/app/faiss_manager.py @@ -1,14 +1,20 @@ -"""Singleton менеджер FAISS GPU индекса. +"""Singleton менеджер FAISS индекса. -Использует IVFFlat (а не HNSW — не поддерживается на GPU). -Поддерживает graceful degradation на CPU если GPU недоступна. +Используется плоский индекс IndexFlatIP, обёрнутый в IndexIDMap2, что даёт: +- отсутствие этапа обучения (в отличие от IVFFlat) — индекс работоспособен сразу, + начиная с первого вектора; +- хранение doc_id прямо внутри индекса (add_with_ids) — не нужен отдельный + маппинг на диске, а поиск сразу возвращает doc_id из PostgreSQL; +- корректную персистентность через faiss.write_index / read_index. + +Векторы модели нормализованы (normalize_embeddings=True), поэтому inner product +эквивалентен косинусной близости. Для масштаба проекта (сотни тысяч — единицы +миллионов документов) полный перебор по FlatIP по скорости приемлем. """ -import json import logging import os from pathlib import Path -from typing import Optional import numpy as np @@ -18,70 +24,83 @@ logger = logging.getLogger(__name__) class FAISSManager: - """Singleton для управления FAISS индексом на GPU.""" + """Singleton для управления FAISS индексом (IndexIDMap2 поверх IndexFlatIP).""" _index = None - _id_map: dict[int, int] = {} # faiss_internal_id -> doc_id (PostgreSQL) - _reverse_map: dict[int, int] = {} # doc_id -> faiss_internal_id + # doc_id -> faiss_id. Для IDMap2 faiss_id == doc_id, но маппинг сохраняем + # для совместимости с вызывающим кодом (plagiarism.embed_documents). + _reverse_map: dict[int, int] = {} _use_gpu: bool = False - _is_trained: bool = False + + @classmethod + def _new_index(cls): + """Создать новый пустой индекс нужного типа.""" + import faiss + + base = faiss.IndexFlatIP(settings.EMBED_DIM) + return faiss.IndexIDMap2(base) @classmethod def load_or_create(cls) -> None: - """ - Загрузить индекс с диска или создать новый. + """Загрузить индекс с диска или создать новый. - Пытается перенести индекс на GPU, при ошибке остаётся на CPU. + Старый несовместимый индекс (например, IVFFlat от предыдущей версии, + который не хранит id_map) безопасно пересоздаётся — полезных векторов в + нём всё равно не было. """ 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 + try: + loaded = faiss.read_index(index_path) + if hasattr(loaded, "id_map"): + cls._index = loaded + cls._rebuild_reverse_map() + logger.info( + f"Загрузка FAISS индекса из {index_path} " + f"({cls._index.ntotal} векторов)" + ) + else: + logger.warning( + "На диске несовместимый FAISS индекс (без id_map) — " + "пересоздаём как IndexIDMap2(IndexFlatIP)" + ) + cls._index = cls._new_index() + cls._reverse_map = {} + except Exception as e: + logger.warning(f"Не удалось загрузить FAISS индекс ({e}) — создаём новый") + cls._index = cls._new_index() + cls._reverse_map = {} 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 + logger.info("Создание нового FAISS индекса IndexIDMap2(IndexFlatIP)...") + cls._index = cls._new_index() + cls._reverse_map = {} + + cls._use_gpu = False # FlatIP на CPU достаточно быстр для целевого масштаба + + @classmethod + def _ensure(cls) -> None: + """Ленивая инициализация индекса при первом обращении.""" + if cls._index is None: + cls.load_or_create() + + @classmethod + def _rebuild_reverse_map(cls) -> None: + """Восстановить _reverse_map из id, хранящихся внутри загруженного индекса.""" + import faiss - # Попытка перенести на GPU try: - res = faiss.StandardGpuResources() - cls._index = faiss.index_cpu_to_gpu(res, 0, cpu_index) - cls._use_gpu = True - logger.info("FAISS индекс размещён на GPU") + ids = faiss.vector_to_array(cls._index.id_map) + cls._reverse_map = {int(i): int(i) for i in ids} 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 + logger.warning(f"Не удалось восстановить reverse_map из индекса: {e}") + cls._reverse_map = {} @classmethod def search(cls, query_vector: np.ndarray, k: int = 20) -> list[tuple[int, float]]: - """ - Поиск k ближайших векторов. + """Поиск k ближайших векторов. Args: query_vector: Нормализованный вектор запроса, форма (768,) @@ -90,21 +109,21 @@ class FAISSManager: Returns: Список кортежей (doc_id, cosine_score), отсортированных по убыванию score """ - if cls._index is None or not cls._is_trained: - logger.warning("FAISS индекс не инициализирован или не обучен, пропускаем поиск") + cls._ensure() + + if cls._index.ntotal == 0: return [] try: query = query_vector.reshape(1, -1).astype(np.float32) - distances, indices = cls._index.search(query, k) + distances, ids = cls._index.search(query, min(k, cls._index.ntotal)) results = [] - for idx, dist in zip(indices[0], distances[0]): + for idx, dist in zip(ids[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))) + # Для IDMap2 idx — это уже doc_id из PostgreSQL + results.append((int(idx), float(dist))) return results @@ -114,95 +133,49 @@ class FAISSManager: @classmethod def add_vectors(cls, vectors: np.ndarray, doc_ids: list[int]) -> None: - """ - Добавить векторы в индекс. + """Добавить (или обновить) векторы в индекс. - Если индекс не обучен и накопилось достаточно векторов — обучить его. + Идемпотентно по doc_id: при повторном эмбеддинге старый вектор документа + удаляется перед добавлением нового, чтобы не плодить дубли. Args: vectors: numpy массив формы (N, 768) doc_ids: Список doc_id из PostgreSQL """ - import faiss + cls._ensure() - if cls._index is None: - cls.load_or_create() + if len(doc_ids) == 0: + return - vectors = vectors.astype(np.float32) + vectors = np.asarray(vectors, dtype=np.float32) + ids = np.asarray(doc_ids, dtype=np.int64) - 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 + # Удалить существующие id, чтобы повторный эмбеддинг не создавал дубли + try: + cls._index.remove_ids(ids) + except Exception: + pass - 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) + cls._index.add_with_ids(vectors, ids) + for doc_id in doc_ids: + cls._reverse_map[int(doc_id)] = int(doc_id) - # Добавить в 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 + """Сохранить индекс на диск.""" + cls._ensure() 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} векторов)") + faiss.write_index(cls._index, str(index_path)) + logger.info(f"FAISS индекс сохранён: {index_path} ({cls._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: