С опытом у RAG-инженера накапливается солидный багаж эвристик и инструментов, которые в определенных задачах превосходят по качеству или скорости стандартные. Фраза «а для этого у меня есть собственный ретривер» звучит с некоторым снобизмом, но добавляет к профессионализму несколько пойнтов.
Хотите в свою коллекцию ретривер, который умеет работать с терминами, плохо различимыми в векторном пространстве эмбеддинга, в частности с именами и названиями? Тогда давайте перейдём от снобизма к практике. Начнём с обработки текста и сегментируем его на фрагменты (чанки). Далее сделаем TFIDF модель, добавим поиск и обернём всё это в ретривер LangChain. Наконец сравним наш ретривер с двумя-тремя стандартными решениями. А Ollama поможет нам с вопросами для бенчмарка.
Всё необходимое для статьи я собрал в репозитарий. В папке data находятся тексты и версия вопросника, в папке gensim - словарь, TFIDF модель, и индекс для поиска, в папке splits сохранён разбитый на фрагменты текст, а в папке retriever – стоп-слова и готовый для импорта модуль.
По традиции я выбрал текст в жанре научной фантастики.
На этот раз - повесть Аркадия и Бориса Стругацких «Жук в муравейнике», 1979 г.
Впервые опубликована в журнале «Знание - сила», в 1979 - 1980 гг..
Ноутбуки туториала:
S1_Processing.ipynb
S2_CreateRetrieverModel.ipynb
S3_CreateCustomRetriever.ipynb
S4_QuestionGeneration.ipynb
S5_RetrieversBenchMark.ipynb
Подготовка (подробно в README.md):
1. Установить Ollama и загрузить модель
2. Создать Python окружение (python==3.11)
3. Установить зависимости из requirements вместе с jupyter для ноутбуков
Шаг первый. Обработка и сегментация текста
Текст романа условно разбит на четыре главы по датам событий, и именованные параграфы. Это авторское разбиение и его следует взять за основу сегментации.
Загружаем текст и маркируем заголовки в стиле маркдаун. Нам нужно три уровня. Первый уровень это название текста (на случай если текстов несколько), второй - название главы, третий - секция параграфа на соответствующую дату. Маркировку сделаем при помощи регулярных выражений, вот основные:
txt = '# ' + txt.strip() # mark title txt = re.sub(r'\n{4}(\d.{3,})\n\n', r'\n\n\n\n## \1\n\n', txt) # mark chapters txt = re.sub(r'\n\n(\«?[А-Я]{1,}.*[А-Я]{2,})\n\n', r'\n\n### \1\n\n', txt) # mark chapter sections
Ещё нужно убрать техническую информацию из текста и лишние пустые строки. Старайтесь ничего не делать руками, всё должно быть в коде. Это особенно важно для корпоративных и юридических документов. Регулярные выражения действуют на весь текст, но к сожалению специфичны. Другой текст потребует настройки выражений. В практике на этом шаге чаще используют различные конверторы в маркдаун, а результат уже корректируется.
Трудно предположить, что разбиение текста сделано Авторами в случайном порыве, если это конечно не художественный приём. Каждый сегмент текста имеет законченный смысл. А раз так, можно использовать MarkdownHeaderTextSplitter из пакета langchain_text_splitters и у нас получится что-то очень близкое к Semantic Chunking. Вот так выглядит сплиттер:
headers_to_split_on = [ ("#", "title"), ("##", "chapter"), ("###", "section"), ] markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on) splits = markdown_splitter.split_text(md_text)
Секции (третий уровень) должны обязательно иметь уникальное имя внутри главы. Если это не так, сплиттер склеит их в один фрагмент. На этот случай я сделал генератор имён на основе части первой строки после маркера. Основная конструкция генератора:
md_text = re.sub(r'\n\n###\n(.{5,25}[ \,\.\-])', r'\n\n### \1\n\1', md_text)
Теперь у нас есть текст, разбитый на чанки в формате LangChain Document. Нужно добавить к метаданным уникальный id (md5 хеш из hashlib от первых 500 символов), размер чанка, название коллекции. В page_content имеет смысл вставить название главы, и оригинальное название параграфа, они пропадаю из текста при сегментации.
Что же у нас получилось?
Рекомендация: если есть возможность, всегда стройте гистограмму длин чанков. Потом пришлёте её девопсу, когда он будет «пытать» вас про размеры контекстного окна.
Итак, у нас есть чанки размером более 35 тысяч символов. Вариант не очень для RAG. Ещё вспомним, что TFIDF предвзято относится к терминам из больших чанков, назначая им больший вес. Поэтому большие чанки нужно аккуратно разбить и отмаркировать третьим уровнем. Говорят, что чанкинг - это искусство.
Для больших чанков я добавил свой сплиттер. Каждый параграф он разбивает на строки и начинает добавлять строки к буферу, проверяя его размер. Если размер превышает параметр TRUNCATE и последняя добавленная строка не является репликой в диалоге, буфер очищается, а к строке добавляется маркер. TRUNCATE нужно подбирать. Ориентируйтесь на качество поиска и кривую распределения, хороший вариант - одна вершина и более менее плавные хвосты. В ноутбуке параметр установлен в 4.000 символов. После разбиения получаем:
Максимальный размер фрагмента около 6к, превышает параметр 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 модель
Для ретривера нужна модель, умеющая в ранжирование. В случае с TFTIDF моделью это три компонента: словарь, модель и индекс для поиска. Но начать придётся снова с обработки текста. Для каждого фрагмента необходимо получить список терминов, при этом убрать стоп-слова, одиночные символы, знаки препинания.
Главный компонент обработки - нормализация терминов, а точнее, приведение терминов к начальной форме. В русском языке это грамматическая форма слова, используемая как основная для лемматизации. Для имен существительных - именительный падеж ед. числа, для прилагательных - именительный падеж ед. числа мужского рода, для глаголов - инфинитив, неопределенная форма.
Для нормализации я выбрал пакет pymorphy3. Он хорош, но не слишком быстрый. 100 чанков обрабатывается за несколько секунд. Вместо pymorphy можно попробовать snowball stemmer для русского языка из пакета NLTK, он работает быстро, но качество ретривера получается несколько хуже. Тут вопрос эксперимента.
Паплайн обработки с нормализатором
morph_analyzer = pymorphy3.MorphAnalyzer() # Filters transform_to_lower = lambda s: s.lower() CLEAN_FILTERS = [ strip_tags, # strip_numeric, # left numbers starting with two digits strip_punctuation, strip_non_alphanum, strip_multiple_whitespaces, transform_to_lower, ] # Filtering of all the unrelevant text elements def cleaning_pipe(text:str) -> list[str]: # Invoking gensim.parsing.preprocess_string method with set of filters processed_words = preprocess_string(text, CLEAN_FILTERS) processed_words = [s for s in processed_words if len(s) > 1] processed_words = [s for s in processed_words if s not in STOP_WORDS] processed_words = [morph_analyzer.parse(s)[0].normal_form for s in processed_words] return processed_words
Сделать модель в gensim несложно, всего несколько строк кода:
dictionary = corpora.Dictionary(processed_docs) bow_corpus = [dictionary.doc2bow(text) for text in processed_docs] model = models.TfidfModel(bow_corpus) index = similarities.SparseMatrixSimilarity(model[bow_corpus], num_features=len(dictionary))
А так работает поисковый механизм:
query_document = cleaning_pipe(query) query_bow = dictionary.doc2bow(query_document) sims = index[model[query_bow]] top_idx = sims.argsort()[-1*N:][::-1]
top_idx – список из N номеров чанков в splits в порядке убывания сходства. Если вы перестроили чанки с другими параметрами или фильтрами, нужно заново сделать словарь, модель и индекс. Иначе индексы разойдутся и ретривер работать не будет. Сходство от нуля до единицы определяется как 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 и выше получится, если отправить в поиск дословный фрагмент текста.
Ещё один пример поиска, similarity 0.429
query = """Кирилл Александров, известный своими антропоморфистскими взглядами, высказал предположение, что саркофаг есть хранилище генофонда Странников. Все известные мне доказательства негуманоидности Странников, заявил он, являются по сути своей косвенными. На самом же деле Странники вполне могут оказаться генетическими двойниками человека. Такое предположение не противоречит ни одному из доступных фактов. Исходя из этого, Александров предлагал все исследования прекратить, вернуть находку в первоначальное состояние и покинуть систему ЕН 9173."""
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 чанков на любой запрос.
Шаблон класса
class RetrieverTemplate(BaseRetriever): """Retriever template""" docs: List[Document] """Documents.""" k: int = 3 """Number of documents to return.""" @classmethod def from_documents( cls, docs: Iterable[Document], **kwargs: Any, ) -> FreqRetriever: """ Create a FreqRetriever instance from a list of langchain Documents. Args: docs: A list of of langchain Documents. **kwargs: Any other arguments to pass to the retriever. Returns: A Retriever instance. """ return cls( docs=docs, **kwargs ) def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> List[Tuple[Document, float]] | None: result_docs = self.docs[:self.k] return result_docs
Загружаем файлы модели, определяем фильтры для обработки запроса, они должны быть идентичны фильтрам для обработки всего набора фрагментов. Далее нужно специфицировать функцию для ранжирования get_top_n. Для простоты я убрал её за границы класса.
Вносим необходимые правки в шаблон. Их немного: добавить параметр with_relevance, чтобы можно было получать оценку similarity, и вставить в функцию get_relevant_documents вызов get_top_n.
Поисковый механизм
def get_top_n(docs: List[Document], query: str, n: int, with_similarity: bool) -> List[Tuple[Document, float]] | List[Document] | None: """Retriever query engine Args: query: text query. n: number of returned documents. Returns: A list of tuples with relevance score or list of Documens. """ query_bow = dictionary.doc2bow(query) sims = similarity_index[retriever_model[query_bow]] qty = sum(sims > 0) if qty > 0: top_idx = sims.argsort()[-1 * n:][::-1] result = [] for idx in top_idx: similarity = round(float(sims[idx]), 3) doc = docs[idx] if with_similarity: result.append((doc, similarity)) else: result.append(doc) return result else: return None
Класс ретривера
class FreqRetriever(BaseRetriever): """TF-IDF retriever based on gensim model""" docs: List[Document] """Documents.""" k: int = 4 """Number of documents to return.""" with_similarity: bool = False """True for return found chunk relevance""" @classmethod def from_documents( cls, docs: Iterable[Document], **kwargs: Any, ) -> FreqRetriever: """ Create a FreqRetriever instance from a list of langchain Documents. Args: docs: A list of of langchain Documents. **kwargs: Any other arguments to pass to the retriever. Returns: A FreqRetriever instance. """ return cls( docs=docs, **kwargs ) def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> List[Tuple[Document, float]] | None: processed_query = cleaning_pipe(query) return_docs = get_top_n(self.docs, processed_query, n=self.k, with_similarity=self.with_similarity) return return_docs
Теперь аккуратно переносим всё необходимое из ноутбука в python-файл, чтобы его можно было импортировать как модуль. Не забудьте про __init__.py
Три момента:
1. Импорт стоп-слов теперь осуществляется из retriever.stop_words
2. Нам не нужно загружать сплиты, они передаются как параметр.
3. В файле freq_retriever.py из репозитария функция get_top_n оформлена как метод класса
Шаг четвёртый. Генерируем вопросы
Ключевая идея на этом шаге — использовать структурированный вывод, главное, чтобы модель умела работать в этом режиме и контекстное окно было больше максимального размера фрагмент + промпт.
Я попробовал три локальные модели:
gemma4:31b - формулировки хорошие, но очень медленно. На один фрагмент текста уходит примерно 5 минут на GPU 16GB.
qwen3:14b — быстро, но коряво в плане русского языка. Можно использовать в крайнем случае.
gemma3:27b-it-qat — компромиссный вариант, примерно 20 секунд на фрагмент. На нём и остановимся.
При помощи обёртки ChatOllama из langchain_ollama объявляем модель:
llm = ChatOllama( model=OLLAMA_MODEL, temperature=0.5, base_url=BASE_URL, seed=42, )
Схема для структурированного вывода. BaseModel из pydantic, а вопросы будут списком.
class QuestionGenerationSchema(BaseModel): """Question generation schema for structured LLM output""" number_of_questions: int questions: list[str]
Собственно генератор:
def query_generation(system_prompt: str, md_text: str) -> QuestionGenerationSchema: content = f'\n\nФрагмент текста:\n\n{md_text}' try: messages = [ { "role": "system", "content": system_prompt }, { "role": "user", "content": content }, ] structured_llm = llm.with_structured_output(QuestionGenerationSchema) response = structured_llm.invoke(messages) if not isinstance(response, QuestionGenerationSchema): raise RuntimeError("LLM did not return correct scheme response") return response except Exception as ex: print(f"Question generation error: {ex}")
Промпт:
QUESTION_GEN_PROMPT = """ Твоя задача - придумать от {min_questions} до {max_questions} вопросов, опираясь на фрагмент художественного текста. Формулируй вопросы так как если бы они задавались на викторине. Убедись, что на вопросы можно ответить, опираясь на фрагмент текста. В формулировке вопроса необходимо и очень **важно** использовать ключевые слова из текста: - имена персонажей - названия мест и организаций - упомянутые в тексте предметы и гаджеты. Правила подготовки вопросов: - Вопросы должны быть заданы на русском языке. - Вопросы не должны повторяться. - В вопросе **нельзя** ссылаться на фрагмент текста, **нельзя** использовать формулировки "согласно описанию в тексте", "согласно тексту", "как указано в тексте", "упоминается в фрагменте текста" Вывод должен быть строго в формате JSON: "number_of_questions": int, // Сколько вопросов ты подготовил для данного фрагмента текста. "questions": [str, ..., str] // Список вопросов. """
Через PromptTemplate задаём значение минимального и максимального количества вопросов для одного фрагмента:
prompt_template = PromptTemplate.from_template(QUESTION_GEN_PROMPT) system_prompt = prompt_template.invoke({ 'min_questions': 1, 'max_questions': 4,}).text
Проходим по списку фрагментов из splits. Если размер фрагмента меньше 400 символов, пропускаем шаг (таких всего три чанка). Забираем текст из атрибута page_content и отправляем в генератор. Из структурированного вывода достаём список вопросов и складываем их в pandas DataFrame с указанием id фрагмента.
30 минут и у меня получилось 377 вопросов на 96 фрагментов. Если использовалась qwen3:14b или gemma3:27b-it-qat, обязательно нужно проверить длину вопроса и отфильтровать слишком короткие в одно-два слова, такое иногда встречается. Сами вопросы не редактируем!
Примеры вопросов (gemma3:27b)
- В каком году Максим Каммерер получил задание от Экселенца найти Льва Вячеславовича Абалкина?
- С какой организации, находящейся на Земле, отбыл Лев Вячеславович Абалкин незадолго до своего исчезновения? - Какой предмет, названный Экселенцем "заккурапия", был использован для передачи информации о связях Льва Абалкина?
- Какую инструкцию Экселенц дал Максиму Каммереру относительно его группы и отчётности по делу Абалкина?
Шаг пятый. Сравниваем ретриверы
Главное, как сравниваем. Проходим по списку вопросов. Отправляем вопросы в ретривер. В метаданных, которые возвращаются в ответ, находим id фрагмента и собираем в список в порядке убывания сходства. Определяем, есть ли id, указанный для вопроса, в этом списке и какова его позиция. На основе этой информации считаем метрики — Hit Rate и MRR (Mean reciprocal rank). Во всех тестах будем просить у ретриверов три документа.
Hit Rate показывает как часто ответ попадает в топ. Для одного вопроса он бинарный - нашёл/не нашёл. MRR это комбо качества ранжирования и полноты поиска. Если он низкий, либо документы находятся, но далеко от начала выдачи, либо ретривер часто «промахивается» и не находит ничего полезного. Предположим, у нас выдача не из трёх, а из десятка документов и нужный документ где-то в конце. В этом случае контекстное окно забивается мусором. MRR ещё можно трактовать как вероятность, что LLM получит верный ответ, «взглянув» только на верхнюю часть списка.
В самом бенчмарке нет ничего хитрого. LangChain даёт единый интерфейс для любого ретривера. Переопределили ретривер, дальше - retriever.invoke(query). Импортируем нужный ретривер, создаём экземпляр и определяем параметры:
from retriever.freq_retriever import FreqRetriever from langchain_community.retrievers import BM25Retriever retriever = BM25Retriever.from_documents(splits, k = 3)
Запускаем цикл по списку вопросов и считаем метрики. Упрощенно, такая конструкция:
num_questions += 1 correct_chunk_id = questions.loc[i, 'chunk_id'] result = retrie.invoke(question) chunks_ids = [] for doc in result: chunks_ids.append(doc.metadata.get('id', '')) if any([idx == correct_chunk_id for idx in chunks_ids]): num_correct += 1 rank += 1./(chunks_ids.index(correct_chunk_id) + 1)
Результаты бенчмарка
Custom Retriever | BM25 Retriever | Vector Retriever | |
Hit Rate@3 | 0.854 | 0.569 | 0.763 |
MRR | 0.748 | 0.477 | 0.615 |
Benchmark time | 551 ms | 198 ms | 38.6 s on cpu |
В таблице представлены усредненные данные по трём прогонам с LLM seed 21, 42 и 63.
Эмбеддинг для векторного ретривера: deepvk/USER-bge-m3
Почему такой слабый результат у стандартных решений?
BM25 «из коробки» не умеет приводить слова к единой форме. Там предусмотрена функция обработки текста, но по умолчанию она лишь разбивает текст на слова. А в векторном поиске не справляется эмбеддинг, для него Щекн Итрч Голован - просто набор символов.
Лайфхак как починить BM25:
Сделайте ещё одну версию сплитов, где текст обработан пайплайном из ноутбука с TFIDF моделью, соответственно и вопрос также нужно обрабатывать. В обоих случаях должен получиться текст из терминов, разделенных пробелами. Ретривер будет выдавать релевантный, но бессвязный набор нормализованных терминов, а по id в метаданных вы без труда получите нужный контекст из основного набора.
Если у вас установлен Qdrant, я добавил в ноутбук S5_RetrieversBenchMark.ipynb необходимый код для него. Можно экспериментировать.
Как использовать получившийся ретривер в RAG-системе? Самый простой вариант - EnsembleRetriever из LangChain, он умеет объединять несколько ретриверов. Важно знать, что уровень сходства с запросом, который получается на выходе нашего ретривера, будет примерно в два раза ниже, чем у ретривера с векторной базой. Прямо объединить выдачи и отсортировать по similarity не получится. Иногда подбирают коэффициент-мультипликатор под конкретный набор исходных документов.
Полезные ссылки:
Маркдаун сплиттер в LangChain
Gensim TFIDF модель
Код BM25Retriever из langchain_community
Код TFIDFRetriever из langchain_community, сделан на sklearn
Хорошая статья с примерами и кодом «How to Measure Your RAG System’s Performance: A Simple Guide» на Medium
P.S. Злоупотребление TFIDF-ретривером может привести к неожиданным находкам.
Берегите себя!
