Обработка данных для RAG-систем

Всем привет, меня зовут Андрей Шахов, я Python-разработчик и Lead Backend Developer в wpp.digital^{_{[ссылка удалена мод.]}}. ML-направлением, а точнее LLM в компании я начал заниматься только в конце 2023 года. Сейчас на задачи подобного рода у меня уходит примерно 40% рабочего времени.

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

Судя по контенту на habr, который я успел просмотреть к этому времени, один из самых популярных способов использования LLM – построение RAG-систем. Далее об этом.

Retrieval Augmented Generation (RAG) – это система, которая передает данные в LLM не в сыром виде от пользователя, а обогащает их контекстом. Система позволяет расширить знание модели, без сложного и долгого обучения. В некоторых случаях даже может хватить работы с удаленным API того же OpenAI, соответственно, ваш бекенд будет маленьким и быстрым. Альтернатива — дообучать модель своими данными, но ответы все еще будут неточными из-за размеров модели, и ресурсов на такое обучение нужно очень много.

Сейчас я расскажу, как готовить данные для RAG-систем. Здесь не будет инструкции, как собрать очередную RAG на каком-нибудь GPT — подобных статей на habr много, без труда найдете подходящую именно для вас. Тут про предварительный шаг – обработку данных, а именно мой опыт обработки отдельно взятой закрытой корпоративной вики.

Флоу процесса обработки

Вот так я выстроил первичный процесс обработки данных после выгрузки:

1. Считал файл страницы;

2. Очистил контент от шума;

3. Разделил очищенный контент на части;

4. Превратил части в векторы с помощью embeddings;

5. Загрузил векторы в БД qdrant, снабдив мета-информацией.

Дальнейшие шаги сильно зависят от качества контента, поэтому из всех шагов я расскажу, как я строил работу на 2-м шаге.

Мне было дано множество страниц в Яндекс. Вики. Все оформлены в разных форматах, с разным контентом и разными редакторами в основе. Страницы были как разводящие, то есть хранящие в себе только ссылки на другие страницы, так и с нужными нам данными. Всего у меня получилось выгрузить порядка 130 страниц в формате Markdown.

Почему markdown?

Markdown отдает наиболее "чистый" текст с точки зрения нужной информации. html-формат, который тоже отдается вики, имеет больше шума. Его обработка даже с помощью beautifulsoup занимает много больше времени, чем чистка markdown. Ранее также у меня был опыт парсинга pdf (как одно-, так и двухколоночного формата), все свелось к извлечению именно текста, с полным игнорированием других элементов.

С другими форматами дела обстоят примерно также – везде много шума, везде от него нужно избавиться. Поэтому выбирать формат, в первую очередь, стоит с точки зрения простоты/быстроты приведения его к чистому тексту.

Например, код одной из страниц (другие содержат корпоративную информацию и не могут быть разглашены, а эта отражает наглядно суть) на html:

<ol>  
<li>  
<p data-line="0"><strong>Ctrl + C</strong> - Копировать</p>  
</li>  
<li>  
<p data-line="2"><strong>Ctrl + X</strong> - Вырезать</p>  
</li>  
<li>  
<p data-line="4"><strong>Ctrl + V</strong> - Вставить</p>  
</li>  
<li>  
<p data-line="6"><strong>Ctrl + Z</strong> - Отменить</p>  
</li>  
<li>  
<p data-line="8"><strong>Ctrl + Y</strong> - Повторить</p>  
</li>  
<li>  
<p data-line="10"><strong>Ctrl + A</strong> - Выделить все</p>  
</li>  
<li>  
<p data-line="12"><strong>Ctrl + S</strong> - Сохранить</p>  
</li>  
<li>  
<p data-line="14"><strong>Ctrl + N</strong> - Создать новый документ/окно</p>  
</li>  
<li>  
<p data-line="16"><strong>Ctrl + O</strong> - Открыть файл</p>  
</li>  
<li>  
<p data-line="18"><strong>Ctrl + P</strong> - Печать</p>  
</li>  
</ol>

И теперь та же страница в Markdown:

1. **Ctrl + C** - Копировать  
2. **Ctrl + X** - Вырезать  
3. **Ctrl + V** - Вставить  
4. **Ctrl + Z** - Отменить  
5. **Ctrl + Y** - Повторить  
6. **Ctrl + A** - Выделить все  
7. **Ctrl + S** - Сохранить  
8. **Ctrl + N** - Создать новый документ/окно  
9. **Ctrl + O** - Открыть файл  
10. **Ctrl + P** - Печать  

Пустые страницы

В вики нашлось несколько нужных страниц, но абсолютно пустых. Такие страницы cчитаются "корневыми" — через них строятся семантические связи между страницами, удалять их нельзя, иначе сломается структура. Были страницы, которые становились пустыми в процессе очистки. Такие страницы сразу удалялись при выгрузке, чтобы не тормозить дальнейший ход работы.

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

1. Очистить контент от markdown-вставок изображений с помощью регулярки {%(.*)%};

2. Очистить контент от ссылок в markdown формате регуляркой \[(.*)]\((.*)\);

3. Если элемент содержит менее 4 символов после этого – пропустить его;

4. Если элемент содержит больше 4 символов — это, с большой вероятностью, какой-то разумный текст, оставить.

О минусах

Такой подход не оказался идеальным, были минусы.

Во-первых, все равно регулярно попадались бесполезные куски текста (например, обрывающиеся на середине предложения), которые только увеличивали процент шума, а значит уменьшали вероятность точного ответа со стороны LLM. Во-вторых, некоторые текстовые куски были маленькими, иногда буквально по несколько слов. Они были малоинформативными, а для векторизации как раз мне нужны были максимально информативные части текста, содержащие законченные мысли и предложения.

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

Чуда не случилось — фокус внимания модели наоборот размазывался. Если она вместе с вопросом получала несколько страниц контекста, ответ чаще всего был неверным, грубо говоря вообще про другое. Регулярно получал ответы вида "bot bot bot..", что является повтором начального токена в промпте, а не корректным ответом. Листинг такого решения следующий:

import re  
from pathlib import Path  

from llama_index.schema import Document  
from marko import Markdown  
from marko.block import BlankLine, BlockElement  
from marko.element import Element  
from tqdm import tqdm  

from utils import get_index  


def get_docs() -> list[Document]:  
	md = Markdown()  
	pages = []  
	for file in Path('pages/').iterdir():  
    	content = file.read_text()  
    	document = md.parse(content)  
    	for block in document.children:  
        	for text in get_text(block):  
            	doc = Document(text=text, metadata={'file_path': file.name})  
            	doc.excluded_embed_metadata_keys = ['file_path']  
            	pages.append(doc)  
	return pages  

  
def get_text(element: BlockElement | Element) -> list[str]:  
	if type(element) is BlankLine:  
        return []  

	if type(element.children) is str:  
    	text = str(element.children)  
     	match_insert = re.match(r'{%(.*)%}', text)  
    	match_image = re.match(r'!\[(.*)](.*)', text)  
    	if match_insert or match_image or len(text) < 4:  
            return []  
    	return [text]  
	texts = []  

	for el in element.children:  
    	texts.extend(get_text(el))  
	return texts  


def main():  
	docs = get_docs()  
	index = get_index()  
	for doc in tqdm(docs):  
    	index.insert(doc)  


if __name__ == "__main__":  
	main()

Флоу процесса обработки №2

На следующем этапе я пересмотрел весь процесс обработки контента и сделал следующие изменения:

1. Очистка контента регулярками стала происходить при загрузке — так файлов стало меньше еще на несколько штук;

2. Добавлена новая регулярка замены: [\n]+ заменялось на один перенос \n;

3. Отказался от marko в пользу деления контента через переносы.

В результате этого в запрос к LLM стали попадать только осмысленные части текста, которые вернула БД в ответ на запрос.

Такие изменения помогли решить проблемы с вопросами, которые могут встречаться на разных страницах.

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

Итоговый листинг получился следующим:

import re  
from pathlib import Path  
 
from llama_index.core.schema import Document  
from tqdm import tqdm  
 
from config import settings  
from utils import get_index  
 

def get_docs() -> list[Document]:  
	pages = []  
	for file in Path(settings.pages_path).iterdir():  
    	content = file.read_text()  
    	content = re.sub(r'{%(.*)%}', "", content)  
    	content = re.sub(r'\[(.*)]\((.*)\)', "", content)  
    	content = re.sub('[\n]+', '\n', content)  
    	pages.extend(get_doc_by_text(content, file))  
	return pages  


def get_doc_by_text(content: str, file: Path) -> list[Document]:  
	pages = []  
	for block in content.split('\n'):  
        doc = Document(text=block, metadata={'file_path': file.name})  
    	doc.excluded_embed_metadata_keys = ['file_path']  
    	if not is_doc_valid(doc):  
        	continue
        pages.append(doc)  
	return pages  


def is_doc_valid(doc: Document) -> bool:  
	if type(doc.text) is not str:  
        return False  
	return len(doc.text.strip()) > 4  


def main():  
	docs = get_docs()  
	index = get_index()  
	for doc in tqdm(docs):  
        index.insert(doc)  
 
 
if __name__ == "__main__":  
	main()

К итогам

Из очевидных истин:

1. Чем меньше размер запроса к LLM, тем меньше частота неточных/неверных ответов;

2. Качество ответов зависит также от промпта, поэтому дальше нужно тюнить уже их.

Теперь к моим личным выводам:

• чем проще формат данных, тем меньше потребуется ресурсов для очищения контента от шума;

• подобрать для всех документов единый 100% рабочий рецепт очистики невозможно, но стремиться к нему нужно, чтобы не очищать данные до второго пришествия;

• декомпозиция процесса позволяет находить слабые стороны процесса очистки и оптимизировать/заменять их.

Welcome в комментарии, буду рад конструктивной критике.