В этой статье, используя технику Retrieval-Augmented Generation ("Поисковая расширенная генерация"), мы настроим русскоязычного бота, который будет отвечать на вопросы потенциальных работников для выдуманного свечного завода в городе Градск.

Что такое RAG?

RAG - это техника, повышающая производительность языковых моделей путём предоставления модели контекста вместе с вопросом.

Существуют разные подходы к использованию RAG. В некоторых случаях используют две LLM, в некоторые встраивают классификаторы или проводят поиск контекста по базе документов. Зависит от целей.

Мы возьмём самый простой:

Последовательность действий:

1. Передадим модели информацию о нашем заводе без дополнительного тюнинга;

2. Создадим базу векторов, где будут храниться ембеддинги ранее заданных вопросов (кэш);

3. При обращении к модели, будем проверять, задавались ли ранее похожие вопросы. Если да, то отдаём ранее сгенерированные ответы.

Зачем использовать кэш?

1. Чтобы увеличить скорость ответов для вопросов, которые задавались ранее.

2. Снизить затраты при использовании платных API (GTP-3.5, GPT-4) для ответов на однотипные и повторяющиеся вопросы.

В качестве модели мы будем использовать адаптивные веса русскоязычной "saiga_mistral_7b_lora", которые натренировал Илья Гусев @Takagiрядом с оригинальной моделью - "Mistral-7B-OpenOrca". Saiga_mistral_7b_lora распространяется по лицензии CC BY 4.0.

При наличии достаточно мощной GPU адаптивные веса, и саму модель можно скачать, запустить у себя локально и пользоваться offline.

Импортируем зависимости и загружаем модель с Hugging Face:

import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer, AutoModel
import torch.nn.functional as F
from langchain.prompts import PromptTemplate

adapt_model_name = "IlyaGusev/saiga_mistral_7b_lora"
base_model_name = "Open-Orca/Mistral-7B-OpenOrca"

tokenizer = AutoTokenizer.from_pretrained(
              base_model_name,
              trust_remote_code=True)

tokenizer.pad_token = tokenizer.eos_token
device_map = {"": 0}

model = AutoPeftModelForCausalLM.from_pretrained(
              adapt_model_name,
              device_map=device_map,
              torch_dtype=torch.bfloat16)

Загружаем "sentence-transformers/all-MiniLM-L6-v2" для получения ембеддингов.

# Load model from HuggingFace Hub
sent_tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
sent_model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')

Ниже - функция для получения эмбеддингов. На вход подаём строку, на выходе получаем torch.tensor размерностью (1, 384):

def get_embedding(sentence):
    
    #Mean Pooling - Take attention mask into account for correct averaging
    def _mean_pooling(model_output, attention_mask):
        token_embeddings = model_output[0] #First element of model_output contains all token embeddings
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    # Tokenize sentences
    encoded_input = sent_tokenizer([sentence], padding=True, truncation=True, return_tensors='pt')

    # Compute token embeddings
    with torch.no_grad():
        model_output = sent_model(**encoded_input)

    # Perform pooling
    sentence_embeddings = _mean_pooling(model_output, encoded_input['attention_mask'])

    # Normalize embeddings
    sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)

    return sentence_embeddings

Создаём базу для эмбеддингов. В промышленных версиях приложений используются vector database, мы же для примера будем использовать обычный torch.tensor.

answers = []
emb_database = torch.empty((0, 384), dtype=torch.float32)

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

user: Небольшой свечной завод, расположенный по адресу ул. Ленина 22, находится
в городе Градск. Доехать до завода от центрального автовокзала можно на автобусе
номер 666. Завод специализируется на производстве свечей различных размеров и
форм. Открытая вакансия: Производственный работник - зарплата 150 тысяч рублей в
месяц. Требования: 6 разряд свечника, 25 лет опыта. Претендент на работу должен
заполнить анкету, предоставить резюме, пройти собеседование с представителями
компании. Для работников завода предусмотрены следующие бонусы: Медицинское
страхование, Абонемент в бассейн. Завод активно проводит праздничные мероприятия
и игры. {question}\nbot: Вот ответ на ваш вопрос длиной не более 10 слов:")

Обернём промт в PromptTemplate из библиотеки langchain. Это аналог f-строки, только с возможностью передавать строку, как зависимую переменную, с последующей передачей ей аргумента.

info_prompt_less10 = PromptTemplate.from_template("user: Небольшой свечной завод, расположенный по адресу ул. Ленина 22, находится в городе Градск. Доехать до завода от центрального автовокзала можно на автобусе номер 666. Завод специализируется на производстве свечей различных размеров и форм. Открытая вакансия: Производственный работник - зарплата 150 тысяч рублей в месяц. Требования: 6 разряд свечника, 25 лет опыта. Претендент на работу должен заполнить анкету, предоставить резюме, пройти собеседование с представителями компании. Для работников завода предусмотрены следующие бонусы: Медицинское страхование, Абонемент в бассейн. Завод активно проводит праздничные мероприятия и игры. {question}\nbot: Вот ответ на ваш вопрос длиной не более 10 слов:")

Вот так выглядит пример из документации langchain:

Напишем функцию для генерации ответа моделью и парсинг ответа:

def get_answer(info_prompt, question):
    
    prompt = info_prompt.format(question=question)   
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), 
                            top_p=0.5,
                            temperature=0.3,
                            attention_mask=inputs["attention_mask"],
                            max_new_tokens=50,
                            pad_token_id=tokenizer.eos_token_id,
                            do_sample=True)

    output = tokenizer.decode(outputs[0], skip_special_tokens=True)

    parsed_answer = output.split("Вот ответ на ваш вопрос длиной не более 10 слов:")[1].strip()

    if "bot:" in parsed_answer:
        parsed_answer = parsed_answer.split("bot:")[0].strip()

    return parsed_answer

Пайплайн:

1. Создаём ембеддинг вопроса:

question = "Каких работников вы ищете?" 
emb = get_embedding(question)
emb

output:
tensor([[ 2.5110e-02,  5.1584e-03, -6.1139e-02, -1.9740e-02, -4.8976e-02,
          5.2549e-02,  1.1428e-01,  1.3328e-01, -2.0873e-02, -6.2296e-02,
         -2.0364e-03,  3.7445e-02,  1.4026e-02,  8.9859e-02, -1.1850e-02,
          ...          ...          ...          ...          ...
          2.8177e-02,  5.4864e-02, -1.2764e-02,  3.7095e-02, -5.4692e-02,
          6.5493e-02, -1.1707e-02, -3.8507e-02, -4.5836e-02, -3.2941e-02,
         -1.8714e-02,  1.0664e-02,  4.8238e-02, -7.9246e-02]])

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

def get_cos_sim(question):
    cos_sim = F.cosine_similarity(emb_database, emb, dim=1, eps=1e-8)
    return cos_sim
  
get_cos_sim(question)

Пока что в базе у нас нету эмбеддингов, поэтому вектор с коэффициентами Отиаи пуст:

output:
tensor([])

3. Получаем ответ от модели, добавляем ембеддинг вопроса в базу ембеддингов, а ответ - в массив answers.

answer = get_answer(info_prompt_less10, question)
emb_database = torch.cat((emb_database, emb), 0)
answers.append(answer)
print(f'Answer from model: {answer}')

output:
Answer from model: Производственный работник.

Как видим модель отвечает согласно переданному контексту. Вот так выглядит кусок в промте: "Открытая вакансия: Производственный работник - зарплата 150 тысяч рублей в
месяц."

Зададим ещё один вопрос и посмотрим коэффициент Отиаи:

question = "Где находится ваш завод?" 
emb = get_embedding(question)
cos_sim = get_cos_sim(question)
cos_sim

output:
tensor([0.5328])

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

Сгенерируем ответ:

answer = get_answer(info_prompt_less10, question)
emb_database = torch.cat((emb_database, emb), 0)
answers.append(answer)
print(f'Answer from model: {answer}')

output:
Answer from model: Завод находится по адресу ул. Ленина 22, Градск.

Ну и последний пример перед тем, как запустить цикл. Зададим вопрос, похожий на предыдущий и посчитаем коэффициент Отиаи:

question = "Где расположен ваш завод?" 
emb = get_embedding(question)
cos_sim = get_cos_sim(question)
cos_sim

output
tensor([0.5483, 0.8426])

Здесь мы уже видим, что один из ранее задаваемых вопрос похож на новый вопрос.

print(f'{cos_sim} {answers=}')

output:
tensor([0.5483, 0.8426]) answers=['Производственный работник.', 'Завод находится по адресу ул. Ленина 22, Градск.']

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

questions = [
"Какой адрес вашего завода?",
"Где находится ваш завод?",
"Какое местоположение вашего завода?",
"Где точно находится ваш завод?",
"Как добраться до вашего завода?",
"Как мне добраться до вашего завода?",
"На каком автобусе добраться до вашего завода?",
"Какие автобусы едут до вашего завода?",
"Что производят на вашем заводе?",
"Какая продукция производится на вашем заводе?",
"Какие товары производятся на вашем заводе?",
"Что именно производится на вашем заводе?",
"Какие изделия производятся на вашем заводе?",
"Какую продукцию я могу найти на вашем заводе?",
"Какик работники вам нужны?",
"Какие должности вы ищете?",
"Какие работники вам нужны в данный момент?",
"Какие специалисты вам требуются?",
"Какие специалисты вам необходимы?",
"Какими навыками и умениями должен обладать специалист?",
"Какие компетенции и знания необходимы специалисту?",
"Какими навыками и качествами должен владеть специалист для успешного выполнения работы?",
"Какие специфические умения и опыт требуются у специалиста?",
"Сколько платят на этой должности?",
"Какова зарплата для этой должности?",
"Какая заработная плата предлагается на этой позиции?",
"Какие условия оплаты труда установлены для этой работы?",
"Какие варианты и размеры заработной платы доступны для этой должности?",
"Какая зарплата ожидается для новых сотрудников на этой позиции?",
"Как устроиться?",
"Как можно устроиться к вам на работу?",
"Какой процесс приема?",
"Какой процесс приема на работу на вашем заводе?",
"Как я могу подать заявку на работу на вашем заводе?",
"Какие шаги мне нужно предпринять, чтобы устроиться к вам на завод?",
"Какие документы и информацию мне необходимо предоставить для трудоустройства на вашем заводе?",
"Какие дополнительные льготы и привилегии предоставляются?",
"Какие бонусы и вознаграждения доступны для сотрудников?",
"Какие дополнительные преимущества и поощрения предусмотрены?",
"Какие бонусные программы и возможности для получения дополнительных выгод имеются?",
"Какие дополнительные бонусы и компенсации предлагаются в рамках трудовых условий?"
]

for q in questions:
    print(q)
    emb = get_embedding(q)
    cos_sim = get_cos_sim(q)
    max_value, max_index = torch.max(get_cos_sim(q), dim=0)

    if max_value > 0.83:
        answer = answers[max_index]
        print(f'DATABASE: {answer}')
    else:
        answer = get_answer(info_prompt_less10, q)
        emb_database = torch.cat((emb_database, emb), 0)
        answers.append(answer)
        print(f'MODEL: {answer}')
    print()

# MODEL и DATABASE сделал заглавными, чтобы визуально было понятно, когда
# отвечает модель, а когда ответ берётся из базы.

output:
Какой адрес вашего завода?
MODEL: ул. Ленина 22, Градск.

Где находится ваш завод?
DATABASE: Завод находится по адресу ул. Ленина 22, Градск.

Какое местоположение вашего завода?
MODEL: ул. Ленина 22, Градск.

Где точно находится ваш завод?
DATABASE: Завод находится по адресу ул. Ленина 22, Градск.

Как добраться до вашего завода?
MODEL: "Автобус номер 666".

Как мне добраться до вашего завода?
DATABASE: "Автобус номер 666".

На каком автобусе добраться до вашего завода?
DATABASE: "Автобус номер 666".

Какие автобусы едут до вашего завода?
DATABASE: ул. Ленина 22, Градск.

Что производят на вашем заводе?
MODEL: Свечи различных размеров и форм.

Какая продукция производится на вашем заводе?
DATABASE: Свечи различных размеров и форм.

Какие товары производятся на вашем заводе?
DATABASE: Свечи различных размеров и форм.

Что именно производится на вашем заводе?
DATABASE: Свечи различных размеров и форм.

Какие изделия производятся на вашем заводе?
DATABASE: Свечи различных размеров и форм.

Какую продукцию я могу найти на вашем заводе?
MODEL: свечи разных размеров и форм.

Какик работники вам нужны?
MODEL: Свечники 6 разряда, 25 лет опыта.

Какие должности вы ищете?
DATABASE: Производственный работник.

Какие работники вам нужны в данный момент?
DATABASE: Свечники 6 разряда, 25 лет опыта.

Какие специалисты вам требуются?
MODEL: Свечник, механик, электрик, контролер.

Какие специалисты вам необходимы?
DATABASE: Свечник, механик, электрик, контролер.

Какими навыками и умениями должен обладать специалист?
MODEL: 6 разряд свечника, 25 лет опыта.

Какие компетенции и знания необходимы специалисту?
DATABASE: 6 разряд свечника, 25 лет опыта.

Какими навыками и качествами должен владеть специалист?
DATABASE: 6 разряд свечника, 25 лет опыта.

Какие специфические умения и опыт требуются у специалиста?
MODEL: 6 разряд свечника, 25 лет опыта.

Сколько платят на этой должности?
MODEL: 150 тысяч рублей в месяц.

Какова зарплата для этой должности?
MODEL: 150 тысяч рублей в месяц.

Какая заработная плата предлагается?
MODEL: 150 тысяч рублей в месяц.

Какие условия оплаты?
MODEL: 150 тысяч в месяц.

Какие варианты и размеры заработной платы доступны для этой должности?
MODEL: Заработная плата 150 тысяч рублей в месяц.

Какая зарплата ожидается для новых сотрудников на этой позиции?
DATABASE: 150 тысяч рублей в месяц.

Как можно устроиться к вам на работу?
MODEL: Заполните анкету и предоставьте резюме.

Как устроиться к вам на работу?Какой процесс приема на работу?
DATABASE: Заполните анкету и предоставьте резюме.

Как я могу подать заявку на работу на вашем заводе?
DATABASE: свечи разных размеров и форм.

Какие шаги мне нужно предпринять, чтобы устроиться к вам на завод?
MODEL: Заполните анкету, предоставьте резюме, пройдите собеседование.

Что сделать, чтобы устроиться к вам на завод?
MODEL: Заполните анкету и предоставьте резюме.

Какие дополнительные бонусы и привилегии предоставляются?
MODEL: Медицинское страхование, абонемент в бассейн.

Какие бонусы и вознаграждения доступны для сотрудников?
MODEL: медицинское страхование, абонемент в бассейн.

Какие бонусные программы предоставляются?
MODEL: медицинское страхование, абонемент в бассейн.

Какие дополнительные бонусы и предлагаются?
DATABASE: Медицинское страхование, абонемент в бассейн.

Из вывода выше мы видим, насколько часто DATABASE отдаёт ответы вместо модели. Сейчас там встречаются ошибки, но это можно поправить, увеличив порог коэффициента Отиаи и подправив промт-контекст.

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

Репозиторий с кодом на GitHub

Google Colab, где можно запустить код на бесплатной Т4 - RAG_LLM_with_cache