Как Нейроюрист ищет по миллионам юридических документов с помощью векторного поиска YDB

Привет, Хабр! Меня зовут Александр Зевайкин, и мы с командой делаем YDB (СУБД Яндекса). В конце прошлого года Яндекс представил специализированного ИИ‑помощника — Нейроюриста, для которого обучил языковую модель на основе Alice AI LLM. Сервис работает на базе RAG, под капотом у которого находится YDB c миллионами различных юридических документов.

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

Как работает векторный поиск с индексом

Получив запрос от пользователя, Нейроюрист ищет в базе подходящие по смыслу документы и добавляет их в промпт, на основе которого языковая модель формулирует ответ. Такой механизм известен как retrieval augmented generation (RAG), а используемый поиск называется векторным поиском.

С помощью индекса векторный поиск можно сделать быстрым. А настройки индекса позволяют всегда находить нужное для RAG количество документов. Но Нейроюрист использует фильтры, и одинаковые настройки индекса применяются к выборкам с количеством документов, отличающимся на порядки.

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

Вообще, векторный поиск можно проводить и без индекса. Сначала нужно вычислить вектор‑эмбеддинг для запроса пользователя и векторные расстояния между ним и эмбеддингами всех текстов в базе. После чего — использовать нужное количество наиболее релевантных вектору запроса текстов. Для нашего RAG мы используем топ-50, чтобы у модели всегда было достаточно контекста для ответа.

У точного векторного поиска без индекса есть всего один недостаток: чем больше эмбеддингов в базе, тем больше времени занимает такой поиск.

Когда количество данных увеличиваетcя, для быстрого поиска можно использовать разные подходы. Например, квантование или векторный индекс. Векторный индекс обеспечивает логарифмическое время поиска за счёт небольшого падения полноты, поэтому разумно использовать именно его. Найденные векторы будут релевантны вектору запроса, но могут быть не самыми близкими, что допустимо для подавляющего большинства ситуаций, особенно для RAG. Такой поиск называют приближённым.

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

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

Главное достоинство приближённого векторного поиска — скорость работы. Для поиска в миллиардах векторов достаточно:

• найти ближайший кластер среди нескольких сотен;

• затем повторить такой поиск несколько раз, спускаясь вниз по уровням дерева;

• и, наконец, найти нужное количество ближайших векторов среди тех нескольких сотен, которые входят в последний найденный кластер.

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

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

Нейроюрист использует фильтры по несбалансированным юридическим данным, поэтому нельзя выбрать хорошие настройки индекса и обеспечить нужную полноту поиска. Настройки задаются на индекс целиком, и для каких‑то значений фильтра в последнем кластере будет слишком много векторов, что замедлит поиск. А для других значений в последнем кластере будет недостаточно векторов, что сделает поиск неполным.

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

Как находить не меньше векторов, чем нужно для RAG

Если объём данных заранее известен, то можно подобрать количество уровней в индексе (за них отвечает параметр levels) и количество кластеров на уровне (за них отвечает параметр clusters).

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

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

Оба этих способа хорошо работают без фильтров или если данные для фильтрации сбалансированы.

Как выглядят данные Нейроюриста?

Особенности юриспруденции: несбалансированное количество данных

Нейроюрист работает с обширной базой юридических данных: Конституцией, федеральными законами, подзаконными актами, приказами ведомств, письмами министерств и федеральных служб, постановлениями и распоряжениями Правительства РФ.

Все эти данные разделены на сферы, например «трудовое право» или «защита прав потребителей». А каждая сфера, в свою очередь, разделена на категории: законодательство, кодексы, судебная практика, комментарии юристов и всё остальное.

Задавая Нейроюристу вопросы, клиенты хотят выбирать сферу права и категорию:

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

Фильтрующий векторный индекс и его одинаковые настройки для всех фильтров

Как отфильтровать результаты векторного поиска с использованием индекса? Фильтровать до поиска нельзя: дерево векторного индекса строится один раз для всех данных, и его нельзя использовать для произвольного набора строк базы данных. Если же фильтровать после поиска, то подходящих результатов может остаться меньше, чем нужно, и качество поиска упадёт.

В YDB фильтрация интегрирована прямо в структуру и алгоритм векторного индекса, так что неподходящие по условиям объекты сразу исключаются из рассмотрения. Как это сделано?

Индекс делится на несколько частей. Вначале строится вторичный векторный индекс для нужных колонок. В случае Нейроюриста это колонки «Сфера права» и «Категория». После чего для каждого уникального элемента вторичного индекса строится отдельное дерево векторного индекса.

В Нейроюристе на момент написания этой статьи 7 сфер права и 5 категорий, так что будет построено 35 деревьев. При поиске с фильтрацией вначале будет выбрано одно из этих деревьев, после чего в выбранном дереве будут найдены нужные для RAG топ-50 векторов, ближайших к вектору запроса. Каждый найденный вектор соответствует какому‑то документу. Нейроюрист забирает эти документы из базы, добавляет в промпт и даёт языковой модели нужный для работы контекст.

Фильтрующий векторный индекс строит все деревья с одинаковыми настройками количества уровней и кластеров. Если выбрать усреднённые параметры в 3 уровня и 40 кластеров на уровень, то даже при расширении полноты поиска (с помощью KMeansTreeSearchTopSize) не для каждой комбинации сферы права и категории удастся найти 50 векторов:

Семейство векторных индексов

Команда разработки Нейроюриста воспользовалась тем, что в 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 с 1 уровнем в дереве и 32 кластерами на этом уровне. А в комбинации «защита прав потребителей» и «судебная практика» 2 506 787 векторов. Поэтому будет использован индекс xxxlarge с 2 уровнями и 128 кластерами:

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

Что получилось

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

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

Семейство векторных индексов позволяет быстро и качественно находить нужное для RAG количество векторов в выборках, отличающихся по объёму на порядки: от тысяч до миллионов векторов:

• время векторного поиска топ-50 векторов на стороне YDB — около 15 мс для p50 и до 200 мс для p99;

• качество поиска по топ-50 кандидатов сохраняется и на больших, и на маленьких подвыборках.

YDB (СУБД Яндекса) доступна как опенсорс‑проект и как коммерческая сборка с открытым ядром. Вы можете запустить её на своих серверах или воспользоваться нашим managed‑решением в Yandex Cloud.