fix(gpu): рабочий FAISS-индекс вместо необучаемого IVFFlat

Семантический поиск (уровень 3) не работал вообще: индекс IndexIVFFlat с
nlist=1024 требует ~40000 векторов для обучения, а до обучения векторы
складывались во временный in-memory _flat_index, который:
  - не сохранялся на диск (save() писал только пустой _index) → терялся при
    рестарте;
  - не участвовал в поиске (search() искал только в необученном _index и
    сразу возвращал []).
Итог: в FAISS всегда было 0 векторов, семантика возвращала пусто.

Заменено на IndexIDMap2(IndexFlatIP): без обучения, работает с первого
вектора, doc_id хранится внутри индекса, корректно персистится. На
нормализованных векторах inner product = cosine, порог 0.75 сохраняет смысл.
Добавлена идемпотентность add_vectors (remove_ids перед add).

Проверено: 66 документов → ntotal=66, поиск возвращает релевантные
результаты со score 0.72-0.78, round-trip save/load сохраняет векторы.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
jze9
2026-07-23 19:28:34 +05:00
parent 9ca2bd1742
commit 38b000703a

View File

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