/rag-finetuned-mistral7b

Primary LanguageJupyter Notebook

Разработка информационной системы ИИ на базе современных LLM-моделей

Решение кейса хакатона DubnaTECH (29 марта - 3 апреля 2024) от комании Нордавинд-Дубна

Целью хакатона была разработка чата с большой языковой моделью.

В качестве базовой модели была выбрана Mistral-7B-Instruct-v0.2 версии GPTQ, которая потребялет относительно мало памяти и подходит для запуска на одном графическом ускорителе. При своем небольшом размере в 7B (7 миллиардов) параметров, её ответы на английском языке сравнимы с моделями ChatGPT ранних версий (175М параметров). Существенным недостатком же является то, что модель плохо оптимизирована для работы с русским языком.

В рамках хакатона качество ответов базовой модели на русском языке было улучшено с помощью fine-tuning'а методом QLoRA на отфильтрованном датасете вопросов-ответов с сервиса Yandex Q.

Fine-tuning делался в среде Google Colab на графическом процессоре NVIDIA A100 в течении 24 минут (60 эпох). График Training и Validation Loss можно будет увидеть ниже. Дообученную LoRA-модель я выложил на Hugging Face.

Поиск по файлам с использованием RAG

Поверх языковой модели реализован поиск по файлам совместно с методом RAG (Retrieval Augmented Generation):

  1. Все файлы в директории documents (точно работает с pdf и txt) индексируются, разбиваются по частям и переводятся в векторное представлением с помощью эмбеддинг-модели.
  2. Запросы пользователей тоже переводятся в векторное представление. Предварительно удаляются стоп-слова русского и английского языков.
  3. Так как мы работаем с векторами, мы можем считать расстояние между ними. Поэтому выбибраем k ближайших векторов к запросу пользователя
  4. Найденные данные добавляются в промт модели как котекст. Например, после слов "Найдено в базе знаний:" идут результаты поиска, а после "Сообщение пользователя:" - сообщение пользователя. Теоретически всё это позволяет управлять поиском используя LLM и естественный язык. Однако сейчас в репозиторий загружены только два файла - мое резюме (чтобы протестировать PDF) и текст романа Мастер и Маргарита, поэтому не удивляйтесь, если вы получите резльутаты связанные с этими файлами. Текст романа очень большой, поэтому, с его текста с большей вероятностью будут выводитьяс при поиске. Чем более подробный запрос, тем больше вероятность, что найдется правльный документ.

Интерфейсом взаимодействия с моделью выступает Telegram-бот. Сейчас есть три команды - /start, /help и /reset. Последняя команда очищает историю диалога (контекст). Бот позволяет работать нескольким пользователям одновременно.

Основые ссылки:

Далее представленны способы запуска моделей: Google Colab и локальный.

Запуск через Google Colab

Ссылка на Colab notebook для запуска бота с моделью: https://colab.research.google.com/drive/1lHbdhq8rJqYoZpMdhFWMj1G1zMlW2qkV?usp=sharing

Модель обучалась и тестировалась в среде Google Collab Pro, из-за необходимости доступа к мощным GPU. Однако даже бесплатные ресурсы Google Colab позволяют запустить телеграм бота, через которого будет происходить общение. В среднем модель потребляет около

  • 6-7 GB GPU RAM
  • 6 GB System RAM.

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

Остальные инструкции есть в самом ноутбуке Google Colab.

Если вы много используете GPU (обычно это NVIDIA T4) в Google Colab, то вам могут на время ограничить доступ к вычислительным ресурсам. Поэтому, если у вас проблема с получением доступа к GPU, попробуйте запустить модель другом аккануте.

Если у вас возникают какие-либо вопросы или проблемы — не стесняйтесь писать мне в телеграм: @examoore

Запуск локально

Этот способ запуска не тестировался, но может сработать для вас. Вам поднадобится:

  • интерпретатор Python 3.10 или больше
  • операционная система Linux
  • графический процессор с поддержкой CUDA 12.2 (другие версии тоже могут работать, но не тестировались) с объемом памяти >8 ГБ
  1. Клонируйте этот репозиторий с помощью команды git clone https://github.com/ggeorg0/rag-finetuned-mistral7b.git
  2. Создайте вирутальное окружение, чтобы не возникало проблем с зависимостями с другими проектами и установите зависимости
python3 -m venv ./.venv && source ./.venv/bin/activate
pip3 install -r requirements.txt
  1. Если вы хотите добавить в систему поиска свои файлы, загрузите их в директорию documents.
  2. Сохраните токен для телеграм-бота в файл telegram_token.txt Такой способ выбран для упрощения процедуры запуска модели. Вы можете использовать следующую команду, для записи токена в файл.
echo "YOUR_TELEGRAM_TOKEN_HERE">telegram_token.txt
  1. Запустите модель с ботом. Первый запуск может занять какое-то время, так как нужно будет скачать модель.
python3 bot.py
  1. Готово! Вы великолепны. Чтобы выйти из виртуального окружения используйте команду deactivate

О датасете и fine-tuning.

Yandex Q - это место, где пользователи могут задавать вопросы и отвечать на них.

Исходный датасет был загружен с Hugging Face

Размер исходных данных: 286 MB: около 664 тысячи строк с различными вопросами и ответами.

Данные был отфильтрованы по размеру ответа (от 400 до 1600 симв), чтобы не отправлять слишком много слов в языковую модель и избежать отдносложных ответов. Были оставльны только те ответы, у которых есть описание.

low_threshould = 400
high_threshould = 1600

clean_dataset = dataset.filter(lambda ds: low_threshould <= len(ds['answer']) <= high_threshould
                                          and len(ds['description']) > 100)

Далее были удалены ответы содержающие обсценную лексику и другие ключевые слова, которые на мой взгляд снижали качество общей выборки. Также удалены начальные предложения с приветсвиями вроде "Здравсвуйте, Семен!"

Таким образом объем данных сократился до 26769 вопросов и ответов, на которых обучалась модель. Посмотреть полный процесс чистки датасета и обучения, можно посмотреть в ноутбуке с Google Colab

Обучение на графическом процессоре NVIDIA A100 в течении 24 минут (60 эпох)

learning_rate: 5e-05
train_batch_size: 8
eval_batch_size: 8

image