About Teletype
Join
H0H1: про HR-аналитику
@h0h1_hr_analytics
LLM
February 1

Bayesium – LLM чат-бот. Часть 1. Сбор и подготовка данных

Будучи вдохновлённым статьёй Кита МакНалти [1], я решил развить изложенную в ней идею и создать чат-бота в Telegram, который будет отвечать на вопросы, опираясь исключительно на книги по байесовской статистике.

Если вы следите за моим каналом, то, возможно, помните, что 2023 год был посвящён байесовской статистике. Основное внимание я уделял двум выдающимся книгам: Doing Bayesian Analysis Джона Крушке [2] и несравненной Statistical Rethinking Ричарда МакЭлрита [3] — моей «библии» по статистике.

Однако время идёт, и, хотя эти книги остаются для меня важными, обращаться к ним ежедневно уже не получается. Что-то неизбежно забывается, и на помощь приходят технологии LLM, которым посвящён этот год на канале.

Можно было бы просто использовать любую доступную LLM-модель (например, DeepSeek) для ответов на вопросы, но мне важно получать ответы исключительно на основе этих двух книг, да ещё и с указанием глав и источников. Именно такую возможность предоставляет фреймворк RAG (Retrieval Augmented Generation) [4], который позволяет создавать контекст для генерации ответов на основе заданных материалов.

В оригинальной статье Кит МакНалт использовал свою книгу, размещённую на GitHub. В моём случае источниками являются PDF-файлы, поэтому в первой части этой серии мы разберёмся, как обработать их и загрузить в базу данных.

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

McElreath R. Statistical rethinking a bayesian course second edition

Перед запуском кода убедитесь, что все книги находятся в рабочей директории. Далее открываем Python и приступаем к работе.

# Устанавливаем необходимые библиотеки (если они ещё не установлены):
# pip install pymupdf
# pip install langchain-community
# pip install chromadb
# pip install sentence-transformers
# pip install --upgrade torch torchvision torchaudio

import fitz               # PyMuPDF для работы с PDF
import pandas as pd       # Для работы с датафреймами
import re                 # Для работы с регулярными выражениями

from langchain_community.document_loaders import DataFrameLoader  # Загрузчик для работы с DataFrame
from langchain.text_splitter import RecursiveCharacterTextSplitter   # Разбиватель текста на части

import chromadb            # Для работы с базой данных векторных эмбеддингов
from chromadb.utils import embedding_functions
from chromadb.utils.batch_utils import create_batches

import uuid                # Для генерации уникальных идентификаторов

Перечислим все названия книг. Последний файл – это решения к книге Джона Крушке.

rethinking_directory = "McElreath R. Statistical rethinking a bayesian course second edition.pdf" 
doingbayesianda_directory = "Krushcke J.K. DoingBayesian DataAnalysis.pdf" 
doingbayesianda_solutions_directory = "Krushcke J.K. DoingBayesian DataAnalysis. Solutions.pdf"

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

def read_pdf(pdf_path):
    """
    Читает PDF-файл и извлекает текст по главам на основе закладок.
    
    Аргументы:
    - pdf_path: путь к PDF-файлу.
    
    Возвращает:
    - main_chapters: словарь, где ключом является название главы, а значением – текст этой главы.
    """
    
    doc = fitz.open(pdf_path)         # Открываем PDF-документ
    bookmarks = doc.get_toc()           # Получаем оглавление (закладки)
    
    main_chapters = {}                # Словарь для хранения глав
    current_main_chapter = None       # Текущая основная глава
    
    # Проходим по всем закладкам
    for i, bookmark in enumerate(bookmarks):
        level, title, page_start = bookmark 
        page_start -= 1  # Переводим страницу в 0-индексацию
        
        # Определяем конец главы: до следующей закладки или до конца документа
        page_end = bookmarks[i + 1][2] - 1 if i + 1 < len(bookmarks) else len(doc)
        
        chapter_text = ""
        # Собираем текст со страниц от начала до конца главы
        for page_num in range(page_start, page_end):
            chapter_text += doc.load_page(page_num).get_text()
        
        # Если уровень закладки равен 1, это основная глава
        if level == 1:
            current_main_chapter = title
            main_chapters[current_main_chapter] = chapter_text.strip()
        # Если это подпункт, добавляем его к текущей главе
        elif current_main_chapter:
            main_chapters[current_main_chapter] += "\n\n" + chapter_text.strip()
    
    return main_chapters

Сначала прочитаем книгу и исключим все не относящиеся к содержательному материалу главы (например, титульные страницы, оглавление, библиографию и т.д.).

# Читаем книгу Ричарда МакЭлрита
rethinking_raw = read_pdf(rethinking_directory)

# Определяем названия глав, которые необходимо исключить
rethinking_chapters_exclusion = [
    'Cover', 
    'Half Title', 
    'Title Page', 
    'Copyright Page', 
    'Table of Contents', 
    'Preface to The Second Edition', 
    'Preface', 
    'Endnotes', 
    'Bibliography', 
    'Citation Index', 
    'Topic Index'
]

# Фильтруем главы, оставляя только содержательные
rethinking_raw = {title: text for title, text in rethinking_raw.items() if title not in rethinking_chapters_exclusion}

# Записываем результаты в DataFrame для дальнейшей обработки
rethinking_df = pd.DataFrame(list(rethinking_raw.items()), columns=["chapter", "text"])
rethinking_df['source'] = 'McElreath R.'  # Добавляем информацию об авторе

# Выводим DataFrame для проверки результата
rethinking_df

Теперь повторим аналогичную процедуру для книги Джона Крушке.

# Читаем книгу Джона Крушке
doingbayesianda_raw = read_pdf(doingbayesianda_directory)

# Определяем названия глав, которые необходимо исключить
doingbayesianda_chapters_exclusion = [
    "1. Front-Matter_2015_Doing-Bayesian-Data-Analysis-Second-Edition-",
    "2. Copyright_2015_Doing-Bayesian-Data-Analysis-Second-Edition-",
    "3. Dedication_2015_Doing-Bayesian-Data-Analysis-Second-Edition-",
    "32. Bibliography_2015_Doing-Bayesian-Data-Analysis-Second-Edition-",
    "33. Index_2015_Doing-Bayesian-Data-Analysis-Second-Edition-"
]

# Фильтруем главы
doingbayesianda_raw = {
    title: text 
    for title, text in doingbayesianda_raw.items() 
    if title not in doingbayesianda_chapters_exclusion
}

# Записываем результаты в DataFrame
doingbayesianda_df = pd.DataFrame(list(doingbayesianda_raw.items()), columns=["chapter", "text"])
doingbayesianda_df['source'] = 'Krushcke J.K.' # Добавляем информацию об авторе

# Выводим DataFrame для проверки результата
doingbayesianda_df

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

def pretty_chapter_title(input_string):
    """
    Приводит название главы к более читабельному виду.
    
    Аргументы:
    - input_string: исходная строка с названием главы.
    
    Возвращает:
    - cleaned_string: очищенная и форматированная строка.
    """
    # Удаляем ведущие цифры и точку (например, "1. ")
    cleaned_string = re.sub(r"^\d+\.\s*", "", input_string)
    # Удаляем всё, что идёт после символа подчеркивания (например, "_2015_...")
    cleaned_string = re.sub(r'_.*#39;, '', cleaned_string)
    # Заменяем дефисы на пробелы
    cleaned_string = cleaned_string.replace("-", " ")
    # Добавляем двоеточие после слова "Chapter" и номера
    cleaned_string = re.sub(r"Chapter (\d+)", r"Chapter \1:", cleaned_string)
    # Исправляем возможное отсутствие апострофа (например, "John s" -> "John's")
    cleaned_string = re.sub(r"(\w) s ", r"\1's ", cleaned_string)
    return cleaned_string.strip()

# Применяем функцию к столбцу с названиями глав в DataFrame книги Крушке
doingbayesianda_df['chapter'] = doingbayesianda_df['chapter'].apply(pretty_chapter_title)

# Выводим DataFrame для проверки результата
doingbayesianda_df

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

def read_pdf_by_pages(pdf_path, chapter_ranges):
    """
    Читает PDF-файл и извлекает текст по заданным диапазонам страниц для каждого раздела.
    
    Аргументы:
    - pdf_path: путь к PDF-файлу.
    - chapter_ranges: список кортежей, где каждый кортеж содержит:
        (название раздела, номер стартовой страницы, номер конечной страницы).
        Нумерация страниц начинается с 1.
    
    Возвращает:
    - main_chapters: словарь, где ключ — название раздела, а значение — текст, извлечённый из указанного диапазона страниц.
    """
    
    # Открываем PDF-документ с помощью PyMuPDF (fitz)
    doc = fitz.open(pdf_path)
    main_chapters = {}  # Словарь для хранения результатов
    
    # Проходим по каждому разделу, определённому в chapter_ranges
    for chapter_title, start_page, end_page in chapter_ranges:
        chapter_text = ""  # Инициализируем переменную для накопления текста раздела
        # Обратите внимание: страницы в PyMuPDF индексируются с 0, поэтому вычитаем 1 из start_page.
        for page_num in range(start_page - 1, end_page):
            # Получаем текст со страницы, удаляем лишние пробелы и добавляем символ перевода строки для разделения страниц
            text = doc.load_page(page_num).get_text()
            chapter_text += text.strip() + "\n"
        # Сохраняем текст раздела в словаре с ключом chapter_title
        main_chapters[chapter_title] = chapter_text.strip()
    
    return main_chapters

# Задаём диапазоны страниц для каждой главы
chapter_ranges = [
    ("Chapter 2. Solutions", 5, 6),
    ("Chapter 3. Solutions", 7, 9),
    ("Chapter 4. Solutions", 10, 16),
    ("Chapter 5. Solutions", 17, 31),
    ("Chapter 6. Solutions", 32, 41),
    ("Chapter 7. Solutions", 42, 51),
    ("Chapter 8. Solutions", 52, 58),
    ("Chapter 9. Solutions", 59, 70),
    ("Chapter 10. Solutions", 71, 80),
    ("Chapter 11. Solutions", 81, 86),
    ("Chapter 12. Solutions", 87, 94),
    ("Chapter 13. Solutions", 95, 104),
    ("Chapter 14. Solutions", 105, 110),
    ("Chapter 15. Solutions", 111, 113),
    ("Chapter 16. Solutions", 114, 122),
    ("Chapter 17. Solutions", 123, 131),
    ("Chapter 18. Solutions", 132, 141)
]

# Читаем PDF-файл с решениями к книге Джона Крушке, используя указанные диапазоны страниц
doingbayesianda_solutions_raw = read_pdf_by_pages(doingbayesianda_solutions_directory, chapter_ranges)

# Преобразуем полученный словарь в DataFrame для удобного анализа и дальнейшей обработки
doingbayesianda_solutions_df = pd.DataFrame(list(doingbayesianda_solutions_raw.items()), columns=["chapter", "text"])
doingbayesianda_solutions_df['source'] = 'Krushcke J.K.' # Добавляем информацию об авторе

# Выводим DataFrame для проверки результата
doingbayesianda_solutions_df

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

# Определяем регулярное выражение для удаления колонтитулов,
# в котором содержатся фрагменты: название книги, год, нумерация страниц и номера глав.
footnote_pattern = r"(Solutions to Exercises in Doing Bayesian Data Analysis 2nd Ed\. by John K\. Kruschke © 2015\.\s*Page \d+ of \d+\s*Chapter(s)? \d+)"

# Удаляем найденные шаблоны из столбца 'text'
doingbayesianda_solutions_df["text"] = doingbayesianda_solutions_df["text"].str.replace(footnote_pattern, "", regex=True)

Соединяем все три документа в один датафрейм.

books = pd.concat([rethinking_df, doingbayesianda_df, doingbayesianda_solutions_df], ignore_index=True)

Далее мы разворачиваем локальную базу данных для наших книг, в которую добавляем документы батчами и используем косинусную меру для оценки схожести эмбеддингов, что впоследствии повлияет на корректность ответов нашего чат-бота. В этой части мы повторяем код из оригинальной статьи Кита МакНалти [1] лишь с небольшими видоизменениями.

# Загружаем данные из объединённого DataFrame с нашими книгами.
loader = DataFrameLoader(books, page_content_column="text")
data = loader.load()

# Разбиваем документы на чанки для лучшей генерации эмбеддингов и поиска релевантного контекста
chunk_overlap — перекрытие между чанками.
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
docs = text_splitter.split_documents(data)
len(docs)
# 4135

# Определяем путь для хранения локальной базы данных Chroma.
CHROMA_DATA_PATH = "chroma_data_bayesium/"
# Используем модель для генерации эмбеддингов.
EMBED_MODEL = "all-MiniLM-L6-v2"


# Создаем клиент ChromaDB, чтобы данные сохранялись локально.
chromadb_client = chromadb.PersistentClient(path=CHROMA_DATA_PATH)

# Задаем имя коллекции, в которой будут храниться текстовые чанки наших книг.
COLLECTION_NAME = "bayesium"

# Определяем embedding_function на базе модели SentenceTransformer.
embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=EMBED_MODEL
)

# Создаем коллекцию в ChromaDB. Обратите внимание, что в metadata указываем использование косинусной меры, что улучшает точность поиска схожих эмбеддингов.
collection = chromadb_client.create_collection(
    name=COLLECTION_NAME,
    embedding_function=embedding_func,
    metadata={"hnsw:space": "cosine"},
)

# Определяем метаданные 
# В метадате передается информация об авторе, названии главы и индексе чанка.
custom_metadata = [
    {
        'source': doc.metadata.get('source', 'Unknown'),
        'chapter': doc.metadata.get('chapter', 'Unknown'),
        'chunk_index': index
    } 
    for index, doc in enumerate(docs)
]

# Далее мы записываем полученные текстовые чанки в базу данных батчами. 
# Функция create_batches автоматически разбивает данные на группы, которые потом загружаются в коллекцию.
batches = create_batches(
    api=chromadb_client,
    ids=[f"{uuid.uuid4()}" for i in range(len(docs))], 
    documents=[doc.page_content for doc in docs], 
    metadatas=custom_metadata  # Use the customized metadata
)

# Добавляем данные в коллекцию батчами. 
for batch in batches:
    print(f"Adding batch of size {len(batch[0])}")
    collection.add(
        ids=batch[0],     
        documents=batch[3], 
        metadatas=batch[2] 
    )

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

question = "How to use ulam()?"

collection.query(
        query_texts=[question],
        n_results=5
    )

Мы закончили с подготовкой данных, в следующей статье мы научим LLM-модель давать ответы на основе этих трех книг.

Код доступен на моём GitHub.

Ссылки

  1. Keith McNulty. How I Created an AI Version of Myself
  2. Krushcke J.K. DoingBayesian DataAnalysis
  3. McElreath R. Statistical rethinking a bayesian course second edition
  4. RAG (Retrieval Augmented Generation) — простое и понятное объяснение
@alexander.botvin
February 1, 19:36
0 views
1 reaction
0 replies
0 reposts