Bayesium – LLM чат-бот. Часть 2. RAG-фреймворк: как настроить LLM для работы со специализированной базой знаний

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

В предыдущей статье [2] мы создали локальную базу данных с книгами по байесовской статистике. В этой части переходим к основной задаче – настройке LLM с применением RAG (Retrieval-Augmented Generation). Коротко напомню: наша цель – настроить LLM таким образом, чтобы ответы формировались строго на основе книг, которые мы разместили в базе данных.

В оригинальной статье МакНалти LLM разворачивается локально на машине, что требует значительных вычислительных ресурсов. Мы же пойдём другим путём и воспользуемся безсерверными (serverless) и бесплатными моделями. Известными платформами для такой задачи являются Hugging Face [3] и OpenRouter [4]. В нашем случае мы будем работать с OpenRouter.

Стартовая страница https://openrouter.ai/

После регистрации на платформе необходимо создать API-ключ. Важно! Сохраните данные о ключе в отдельном документе – после первого показа вы уже не сможете их просматривать.

Особенность платформы OpenRouter в том, что она позволяет легко отличить бесплатные модели от платных по ключевому слову free. Для чат-бота Bayesium я остановился на модели «Llama 3.1 70 B». Сейчас на OpenRouter бесплатно доступна ультрасовременная китайская модель DeepSeek, но для нашей задачи такая мощность не требуется, а стоять в очереди, которая неизбежна из-за хайпа и высокой нагрузки на DeepSeek, совсем не хочется.

Переходим к настройке RAG в Python.

# Установка необходимых пакетов:
# pip install openai

# Импортируем библиотеку для работы с ChromaDB
import chromadb
from chromadb.utils import embedding_functions
from chromadb.utils.batch_utils import create_batches

# Импортируем клиент OpenAI для работы с API OpenRouter
from openai import OpenAI

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

# Указываем путь к локальной базе данных ChromaDB
CHROMA_DATA_PATH = "chroma_data_bayesium/"
# Используем предобученную модель для получения эмбеддингов
EMBED_MODEL = "all-MiniLM-L6-v2"
# Название коллекции, в которой хранятся документы
COLLECTION_NAME = "bayesium"

# Инициализируем клиента ChromaDB для работы с персистентным хранилищем
chromadb_client = chromadb.PersistentClient(path=CHROMA_DATA_PATH)

# Переинициализируем функцию эмбеддинга с использованием выбранной модели
embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=EMBED_MODEL
)

# Получаем уже созданную коллекцию с книгами по байесовской статистике
collection = chromadb_client.get_collection(
    name=COLLECTION_NAME,
    embedding_function=embedding_func  
)

Создаём клиента OpenAI для доступа к API OpenRouter.

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key="<ВАШ API KEY>", 
)

Ниже представлена функция, которая принимает вопрос, выбирает n наиболее подходящих документов из базы и передаёт их LLM для генерации ответа. При этом в ответе будет содержаться информация о книге и главах (из метаданных), на основе которых был сформирован ответ.

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

def ask_question(question: str, n_docs: int = 3) -> str:
    """
    Функция для обработки вопроса.
    
    Аргументы:
      question (str): Текст вопроса.
      n_docs (int): Количество документов, извлекаемых из базы для формирования ответа.
      
    Возвращает:
      str: Сгенерированный ответ LLM.
    """
    
    # Получаем коллекцию из базы ChromaDB
    collection = chromadb_client.get_collection(
        name=COLLECTION_NAME,
        embedding_function=embedding_func  
    )
    
    # Выполняем поиск по схожести: получаем n_docs наиболее релевантных документов
    results = collection.query(
        query_texts=[question],
        n_results=n_docs
    )
    
    # Формируем контекст с текстом документа и соответствующими метаданными (глава, источник)
    context_with_metadata = ""
    referenced_chapters = set()  # Множество для хранения уникальных названий глав
    referenced_sources = set()   # Множество для хранения уникальных источников

    # Проходим по найденным документам
    for i, document in enumerate(results['documents'][0]):
        # Получаем метаданные для текущего документа
        metadata = results['metadatas'][0][i]  
        chapter = metadata.get('chapter', 'Unknown Chapter')  
        referenced_chapters.add(chapter) 
        source = metadata.get('source') 
        referenced_sources.add(source)
        
        # Формируем строку с информацией о документе
        context_with_metadata += f"Document {i+1} (Chapter: {chapter}):\n"
        context_with_metadata += f"{document}\n"
        context_with_metadata += "-" * 50 + "\n"

    # Объединяем названия глав и источники в строку для дальнейшего упоминания в ответе
    chapters_list = ", ".join(referenced_chapters)
    sources = ", ".join(referenced_sources)

    # Формируем окончательный запрос для LLM, включающий вопрос и контекст
    prompt = f"""
    You are an expert on statistics and its applications to analytics.
    Here is a question: {question}
    Answer it strictly based on the following information:
    {context_with_metadata}
    When providing the answer, reference the chapters used by stating "This information is from {chapters_list} by {sources}".
    If the context does not provide enough information, say "I cannot answer this question based on the provided context."
    """

    # Оборачиваем запрос в формате сообщений для модели
    messages = [
        {"role": "user", "content": prompt}
    ]

    # Отправляем запрос на генерацию ответа LLM с помощью OpenRouter
    completion = client.chat.completions.create(
        model="meta-llama/llama-3.1-70b-instruct:free", 
        messages=messages,
        max_tokens=1000,
        temperature=0.25
    )

    # Возвращаем сгенерированный ответ
    return completion.choices[0].message.content

Протестируем модель вопросом о том, как построить модель в R используя функцию ulam(), эта информация описана в книге МакЭлрита Statistical Rethinking [5].

question = "Give an example of regression using the Ulam function in R."

answer = ask_question(question)
print(answer)

Пример ответа на вопрос

Подражая Киту МакНалти зададим вопрос о том, чего точно нет в содержании этих книг: например, как построить дашборд?

question = "Give an example how to create a dashboard"

answer = ask_question(question)
print(answer)

Пример ответа на вопрос

Мне очень понравился этот ответ, нам не только рассказали, что в книгах такой информации нет, но рекомендовали R, ggplot2 и Shiny. Очень круто!

Лично мне неудобно отправлять запросы через интерфейс Jupyter, поэтому, учитывая, что я много времени провожу в Telegram, в следующей статье я расскажу, как настроить полноценного чат-бота на основе данной модели.

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

Ссылки

Keith McNulty. How I Created an AI Version of Myself
Ботвин А.Ю. Bayesium – LLM чат-бот. Часть 1. Сбор и подготовка данных
https://huggingface.co/
https://openrouter.ai/
McElreath R. Statistical rethinking a bayesian course second edition