RAG: Как собрать свой ретривер для особых случаев

RAG: Как собрать свой ретривер для особых случаев

С опытом у RAG-инженера накапливается набор эвристик и инструментов, которые в определённых задачах превосходят стандартные решения по качеству или скорости. Фраза «у меня есть собственный ретривер» звучит с лёгким снобизмом, но добавляет профессионализму веса.

Хотите создать ретривер, способный работать с терминами, плохо различимыми в векторном пространстве, — например, с именами и названиями? Тогда перейдём от слов к делу. Мы обработаем текст, разобьём его на чанки, построим TF-IDF модель, добавим поиск и обернём всё в ретривер на LangChain. Затем сравним его с несколькими стандартными решениями. Для генерации вопросов воспользуемся Ollama.

Весь код и данные собраны в репозитории. В папке data — исходный текст и версия вопросника, в gensim — словарь, модель TF-IDF и поисковый индекс, в splits — разбитые на фрагменты тексты, в retriever — стоп-слова и готовый модуль для импорта.

В качестве текста выбрана научно-фантастическая повесть Аркадия и Бориса Стругацких «Жук в муравейнике» (1979), впервые опубликованная в журнале «Знание — сила».

Ноутбуки туториала:

  • S1_Processing.ipynb
  • S2_CreateRetrieverModel.ipynb
  • S3_CreateCustomRetriever.ipynb
  • S4_QuestionGeneration.ipynb
  • S5_RetrieversBenchMark.ipynb

Подготовка (см. подробности в README.md):

  • Установить Ollama и загрузить модель
  • Создать Python-окружение (python==3.11)
  • Установить зависимости из requirements, включая jupyter

Шаг первый. Обработка и сегментация текста

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

Текст загружается и размечается в стиле Markdown с тремя уровнями заголовков: название произведения, название главы, название секции (по дате). Разметка выполняется с помощью регулярных выражений.

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

Авторское деление текста, как правило, не случайно. Каждый сегмент имеет законченный смысл. Поэтому можно использовать MarkdownHeaderTextSplitter из langchain_text_splitters — это приближает результат к Semantic Chunking.

Секции (третий уровень) должны иметь уникальные имена внутри главы. В противном случае сплиттер объединит их. Для решения этой проблемы используется генератор имён на основе начальной части строки после маркера.

После сегментации к каждому чанку добавляются метаданные: уникальный id (md5-хеш от первых 500 символов), размер чанка, название коллекции. В page_content вставляются название главы и оригинальное название параграфа — они теряются при сегментации.

Рекомендация: всегда стройте гистограмму длин чанков. Это пригодится при обсуждении с devops-командой ограничений по размеру контекстного окна.

Получились чанки до 35 тысяч символов — не лучший вариант для RAG. Кроме того, TF-IDF предвзят к терминам из длинных чанков, присваивая им больший вес. Поэтому большие фрагменты нужно дополнительно разбить.

Для этого используется кастомный сплиттер: он разбивает параграфы на строки, добавляет их в буфер и проверяет длину. Если длина превышает порог TRUNCATE, а последняя строка не является репликой в диалоге, буфер сбрасывается, и добавляется маркер. Параметр TRUNCATE подбирается эмпирически — ориентируйтесь на распределение длин и качество поиска. В ноутбуке он установлен на 4000 символов.

Максимальный размер чанка — около 6000 символов, что превышает TRUNCATE из-за запрета разбивать диалоги.

Document(metadata={'title': 'ЖУК В МУРАВЕЙНИКЕ', 'chapter': '01.06. — 13.01. СЛОН — СТРАННИКУ.', 'section': 'И тут же я понял еще одну...', 'id': '332c2c4dc2f1141827c79cdab488b978', 'size': 360, 'collection': 'beetle_in_anthill'}, page_content='01.06. — 13.01. СЛОН — СТРАННИКУ.\n\nИ тут же я понял еще одну вещь. Вернее, не понял, а почувствовал. А еще точнее — заподозрил. Вся эта громоздкая папка, все это обилие бумаги, вся эта пожелтевшая писанина ничего не дадут мне, кроме, может быть, еще нескольких имен и огромного количества новых вопросов, опять–таки не имеющих никакого отношения к вопросу КАК.')

В итоге получено 99 чанков — достаточно для бенчмарка. Сохраняем splits в формате pickle и переходим к построению модели.

Шаг второй. Построение модели на Gensim

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

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

Для нормализации используется pymorphy3. Он точен, но не очень быстр — обработка 100 чанков занимает несколько секунд. Альтернатива — snowball stemmer из NLTK, работает быстрее, но качество хуже. Выбор зависит от задачи и требований к скорости.

Создание модели в Gensim занимает несколько строк кода. Поисковый механизм работает так:

top_idx — список индексов чанков в порядке убывания сходства. Если вы меняли параметры сегментации, нужно пересоздать словарь, модель и индекс — иначе индексы не совпадут, и ретривер сломается. Сходство (от 0 до 1) определяется как sims[idx], где idx — индекс из top_idx.

Проверим поиск:

query="""Что напомнило Бромбергу о саркофаге в связи с Александром Дымком?"""

chapter: 4 июня 78–го года
section: Сегодня а вернее вчера...
similarity: 0.343

chapter: 3 июня 78–го года
section: ЗАСТАВА НА РЕКЕ ТЕЛОН
similarity: 0.166

chapter: 4 июня 78–го года
section: ЛЕВ АБАЛКИН У ДОКТОРА БРОМБЕРГА
similarity: 0.139

Сходство 0.2–0.3 — хороший результат для этой модели. Выше 0.5 получается, только если искать дословный фрагмент.

query= """Кирилл Александров, известный своими антропоморфистскими взглядами, высказал предположение, что саркофаг есть хранилище генофонда Странников..."""

chapter: 4 июня 78–го года
section: Кирилл Александров...
similarity: 0.429

chapter: 4 июня 78–го года
section: Что же касается Геннадия...
similarity: 0.086

chapter: 4 июня 78–го года
section: ТАЙНА ЛИЧНОСТИ ЛЬВА АБАЛКИНА
similarity: 0.069

Шаг третий. Создание кастомного ретривера

Чтобы создать ретривер в LangChain, нужно унаследовать класс от BaseRetriever из langchain_core.retrievers. Подготовлен шаблон, возвращающий первые k чанков по запросу.

Загружаются файлы модели, определяются фильтры для обработки запроса — они должны совпадать с теми, что использовались для обработки чанков. Также реализуется функция get_top_n для ранжирования. Для простоты она вынесена за пределы класса.

Вносятся правки в шаблон: добавляется параметр with_relevance для возврата оценки similarity, и в get_relevant_documents вставляется вызов get_top_n.

Далее весь код переносится из ноутбука в Python-файл для импорта как модуля. Не забываем про __init__.py.

Три важных момента:

  • Стоп-слова импортируются из retriever.stop_words
  • Чанки не загружаются внутри класса — они передаются как параметр
  • В файле freq_retriever.py функция get_top_n реализована как метод класса

Шаг четвёртый. Генерация вопросов

Ключевая идея — использовать структурированный вывод. Модель должна поддерживать этот режим, а контекстное окно — быть больше, чем максимальный размер чанка плюс промпт.

Протестированы три локальные модели:

  • gemma4:31b — качественные формулировки, но очень медленно (около 5 минут на фрагмент при GPU 16 ГБ)
  • qwen3:14b — быстро, но с ошибками в русском языке. Подходит в крайнем случае
  • gemma3:27b-it-qat — компромисс: около 20 секунд на фрагмент. Выбираем эту

Для структурированного вывода используется BaseModel из Pydantic, вопросы — в виде списка.

Генератор настраивается через PromptTemplate с указанием минимального и максимального количества вопросов на фрагмент.

Обрабатываем список чанков. Фрагменты короче 400 символов пропускаем (их всего три). Текст из page_content отправляется в генератор. Из структурированного ответа извлекается список вопросов и сохраняется в pandas.DataFrame с указанием id чанка.

Через 30 минут получено 377 вопросов по 96 фрагментам. При использовании qwen3:14b или gemma3:27b-it-qat нужно отфильтровать слишком короткие вопросы (одно-два слова). Сам текст вопросов не редактируется!

  • В каком году Максим Каммерер получил задание от Экселенца найти Льва Вячеславовича Абалкина?
  • С какой организации, находящейся на Земле, отбыл Лев Вячеславович Абалкин незадолго до своего исчезновения?
  • Какой предмет, названный Экселенцем "заккурапия", был использован для передачи информации о связях Льва Абалкина?
  • Какую инструкцию Экселенц дал Максиму Каммереру относительно его группы и отчётности по делу Абалкина?

Шаг пятый. Сравнение ретриверов

Сравнение проводится так: проходим по всем вопросам, отправляем их в ретривер, получаем список чанков с метаданными. Находим id исходного чанка и определяем его позицию в выдаче. На этой основе считаем две метрики: Hit Rate и MRR (Mean Reciprocal Rank). Во всех тестах запрашиваем по три документа.

Hit Rate@3 показывает, как часто правильный ответ попадает в топ-3. Для одного вопроса — бинарный показатель: нашёл/не нашёл.

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

Бенчмарк прост: LangChain предоставляет единый интерфейс для всех ретриверов. Достаточно переопределить retriever и вызвать retriever.invoke(query). Импортируем ретривер, создаём экземпляр, задаём параметры — и запускаем цикл по вопросам.

Результаты бенчмарка

Custom Retriever

BM25 Retriever

Vector Retriever

Hit Rate@3

Benchmark time: 38.6 с на CPU

Данные усреднены по трём запускам с seed 21, 42 и 63. Эмбеддинг для векторного ретривера: deepvk/USER-bge-m3.

Почему стандартные решения показали слабый результат?

  • BM25 «из коробки» не нормализует слова. Он просто разбивает текст на токены, не приводя их к начальной форме.
  • Векторный ретривер не справляется с именами вроде «Щекн Итрч Голован» — для эмбеддинг-модели это просто набор символов.

Лайфхак по улучшению BM25: создайте вторую версию чанков, где текст уже прошёл нормализацию (как в TF-IDF-пайплайне). Запрос тоже нужно обрабатывать тем же пайплайном. Ретривер будет возвращать релевантные, но бессвязные термины — по id из метаданных вы легко найдёте нужный контекст в основном наборе.

Если установлен Qdrant, в ноутбуке S5_RetrieversBenchMark.ipynb есть готовый код для экспериментов.

Как использовать кастомный ретривер в RAG-системе? Самый простой способ — EnsembleRetriever из LangChain, который объединяет несколько ретриверов. Важно: значения similarity у нашего ретривера примерно вдвое ниже, чем у векторного. Просто объединить и отсортировать по релевантности нельзя. Иногда подбирают коэффициент-множитель под конкретный набор документов.

P.S. Злоупотребление TF-IDF-ретривером может привести к неожиданным находкам. Берегите себя!

Читать оригинал