Как Нейроюрист ищет ответ среди миллионов документов: векторный поиск в YDB

Нейроюрист — ИИ-помощник по правовым вопросам, которого Яндекс представил в конце прошлого года. Он работает по схеме RAG: сначала находит в базе документы, релевантные вопросу по смыслу, и только потом языковая модель формулирует ответ, опираясь на найденное. Иначе говоря, прежде чем ответить, система подкрепляет генерацию текста результатами предварительного поиска.

Зачем векторному поиску индекс

Самый простой вариант поиска по смыслу — сравнить вопрос с каждым документом в базе и выбрать наиболее похожие. Для этого вычисляется эмбеддинг запроса и расстояние от него до эмбеддинга каждого документа. Нейроюристу для ответа нужно топ-50 таких документов: этого хватает языковой модели, чтобы сформировать контекст.

У такого способа один существенный недостаток: время растёт линейно с объёмом данных. Для 1 тыс. документов поиск займёт пять миллисекунд. Для 100 тыс. документов — уже 300 миллисекунд. У Нейроюриста документов на порядки больше — счёт идёт на миллионы, и их число растёт каждый день.

Векторный индекс решает эту проблему. Он устроен как дерево кластеров. Сначала все векторы делятся на несколько кластеров первого уровня, и для каждого кластера считается центроид — усреднённый вектор, представляющий кластер целиком. Затем каждый кластер делится на кластеры второго уровня со своими центроидами, и так повторяется столько раз, сколько уровней задано в индексе.

Поиск идёт по этому дереву сверху вниз. Вектор запроса сравнивается с центроидами первого уровня, и выбирается ближайший. Затем — с центроидами второго уровня внутри выбранного кластера, и снова выбирается ближайший. Так продолжается, пока не будет достигнут последний уровень: там вектор запроса сравнивается уже с конкретными векторами внутри найденного кластера, и среди них отбираются ближайшие.

Поэтому поиск по индексу так хорошо масштабируется: чтобы найти нужные векторы среди миллиарда, достаточно один раз выбрать ближайший кластер среди нескольких сотен, повторить это на каждом уровне дерева, а затем перебрать лишь несколько сотен векторов, что остались в последнем кластере. Цена такой скорости — точность: индекс находит достаточно близкие векторы, но не гарантирует, что среди них окажется самый точный. Для RAG этого обычно хватает, и такой поиск называют приближённым.

Проблема несбалансированных данных

Чтобы в последнем кластере оказалось нужное число векторов, индекс нужно настраивать заранее: задать число уровней дерева и число кластеров на каждом уровне — под ожидаемый объём данных.

У Нейроюриста есть фильтры. Вы выбираете сферу права и категорию документа: законодательство, кодексы, судебную практику или комментарии юристов. И здесь возникает сложность — количество документов для разных сочетаний фильтров отличается не в разы, а на порядки. Судебная практика по защите прав потребителей — это 2,5 млн документов. Законодательство по корпоративному праву — только 6 тыс.

Настройки индекса задаются на всё дерево сразу, а не отдельно под каждый фильтр. Если взять усреднённые параметры — три уровня и 40 кластеров на уровень, — для части фильтров поиск найдёт нужные 50 документов, а для других не наберёт и половины:

Фильтрующий индекс и его предел

Фильтровать документы до поиска нельзя: дерево индекса строится один раз для всей таблицы, и под конкретный фильтр его не перестроить на лету. Фильтровать после поиска тоже нельзя: часть подходящих документов к этому моменту уже будет потеряна, и для узких фильтров, таких как корпоративное право, эта потеря окажется особенно заметной.

В YDB фильтрация встроена прямо в индекс. Сначала строится вторичный индекс по колонкам фильтра — у Нейроюриста это сфера права и категория. Затем для каждой уникальной комбинации значений строится собственное дерево векторного индекса. При семи сферах права и пяти категориях получается 35 таких деревьев. Когда приходит запрос с фильтром, сначала выбирается нужное дерево и только внутри него ищутся топ-50 ближайших векторов — документы за пределами фильтра в сравнении вообще не участвуют.

Это решает одну проблему, но не решает другую. Если все 35 деревьев построить с одинаковыми параметрами, разница в объёме данных между ними сохранится: дереву на 6 тыс. документов и дереву на 2,5 млн по-прежнему нужны разные настройки — иначе получится та же таблица, что и выше.

Семейство индексов вместо одного

В YDB можно создать несколько векторных индексов для одной таблицы и явно выбирать нужный в запросе:

PRAGMA ydb.KMeansTreeSearchTopSize = "<значение>";
SELECT
  id,
  metadata,
  Knn::CosineDistance(embedding, $query_vector) as distance
FROM embeddings VIEW <имя_векторного_индекса>
WHERE index_id = $index_id
ORDER BY distance
LIMIT 50;

Команда Нейроюриста воспользовалась этой возможностью и построила для таблицы несколько индексов с разными параметрами. Каждый индекс отвечает за свой диапазон объёма данных. Комбинация «корпоративное право + законодательство» с её 6022 векторами использует индекс tiny: один уровень, 32 кластера. Комбинация «защита прав потребителей + судебная практика» с 2 506 787 векторами использует индекс xxxlarge: два уровня, 128 кластеров на уровень:

Число уровней и кластеров подбиралось так, чтобы в последнем кластере оставалось от 50 до 150 векторов: меньше — и приходится перебирать много кластеров с малым числом векторов в каждом; больше — и приходится перебирать слишком много векторов внутри одного кластера. Оба случая одинаково избыточны.

Результат

Возможность создавать несколько векторных индексов на одну таблицу, ограничивать их предикатами и явно выбирать нужный в запросе позволила команде Нейроюриста собрать семейство индексов: каждый оптимален для своего диапазона объёма данных.

Поиск топ-50 векторов на стороне YDB занимает около 15 миллисекунд на медиане и до 200 миллисекунд на 99-м перцентиле. Эта скорость не зависит от того, идёт ли поиск среди 6 тыс. документов или среди 2,5 млн — для каждого масштаба заранее подготовлен подходящий ему индекс.

YDB доступна как опенсорс-проект, как коммерческая сборка с открытым ядром и как управляемый сервис Yandex Cloud — включая бессерверный режим с Free Tier на разработку и тестирование подобных сценариев.