Приветствую, хабровчане!
Сегодня хочу рассказать вам историю о том, как я обучил простую и компактную независящую от языка (language agnostic) модель-эмбеддер, которая умеет работать с техническими текстами о PHP и способна извлекать схожие эмбеддинги для параллельных текстов на английском и русском языках.
Основная причина, по которой я решил заняться этим проектом, заключается в том, что мои заметки, код и документация, накопленные за более чем десять лет практики, представляют собой солянку текстов о разных технологиях, языках программирования, пометки о настройке серверов Linux и т.д. на русском и английском языках. Поэтому мне захотелось сделать Retrieval-Augmented Generation (RAG) помогалку, которая сможет принимать запросы пользователя (меня) и эффективно находить информацию в столь разношерстой базе данных, независимо от того на каком языке я сделал запрос и на каком языке написана документация.
Для достижения этой цели как-раз и необходима независимая от языка модель-эмбеддер, которая будет одинаково хорошо работать с техническими текстами на русском и английском языках.
Ещё одним важным аспектом было то, чтобы модель потребляла как можно меньше ресурсов и, если возможно, чтобы её можно было преобразовать в формат GGUF.
Но прежде чем приступить к созданию своего собственного велосипеда, я решил поискать готовые решения, ведь подобная идея очевидна и, возможно, уже реализована другими.
Спойлер: идея не нова, и подобных решений уже достаточно много.
Обзор существующих решений
Для построения системы, которая может извлекать одинаковые эмбеддинги для схожих текстов на русском и английском языках, существует несколько решений, например...
Multilingual Universal Sentence Encoder (mUSE)
Ссылки: arxiv:1907.04307 , kaggle, github
Это проект разработан инженерами Google и поддерживает 16 языков.
Свойства: ~110m параметров, принимает на вход 128 токенов текста и извлекает из них 512-мерный эмбеддинг.
Плюс: поддерживает русский язык.
Минусы: модель основана на Tensorflow, а так же что с 2019го года не было обновлений.
Multilingual Unsupervised and Supervised Embeddings (MUSE)
Ссылки: arxiv:1710.04087, github
Это одна из первых попыток инженеров FB создать модель которая способна выполнять задачи по извлечению независящих от языка эмбеддингов.
Плюс: поддерживает русский язык.
Минусы: в наличии имеются веса для пар языков, навроде en-ru, en-de и т.д., весов нет на HuggingFace, ну и с 2018го года проект не развивается.
Language-Agnostic SEntence Representations (LASER)
Ссылки: arvix:2205.12654, github, pypi
Ещё одна модель разработана инженерами FB и, как сказано в ридми на GitHub, поддерживает более 200 языков (хотя если пройти по ссылочкам и посчитать то получится 147 языков).
Свойства: ~256m параметров, принимает 1024 токенов на вход и извлекает из них 1024-мерный эмбеддинг.
Плюсы: она основана на PyTorch и имеет логику переключения между языками которая явно перекочевала из NLLB (о которой я кстати рассказывал в публикации "Перевод на разные языки используя модель NLLB" у себя в блоге на Дзен).
Минусы: весов нет на HuggingFace, а модель несовместима с llama.cpp поэтому её не получится конвертировать в GGUF, чтобы можно было запускать на слабом железе (или же в паре с ollama).
Multilingual Sentence-BERT (SBERT)
Ссылки: arXiv:1908.10084, сайт
Модели Sentence-BERT представляют собой модифицированную версию предобученной BERT, специально адаптированную для генерации эмбеддингов предложений, multilingual версия позволяет извлекать эмбеддинги из текста на разных языках, а paraphrase модели позволяют извлекать похожие эмбеддинги парафраз на разных языках.
Вот пару примечательных моделей, обученных разными способами:
paraphrase-multilingual-MiniLM-L12-v2 имеет 118m параметров, принимает 256 токенов на вход и возвращает 384-мерный эмбеддинг.
paraphrase-multilingual-mpnet-base-v2 имеет 278m параметров, принимает на вход 512 токенов и возвращает 768-мерный эмбеддинг.
Обе эти модели обучены на комбинации из датасетов:
SNLI о котором говорится в публикации "A large annotated corpus for learning natural language inference" (570k примеров)
Multi-Genre NLI, подробнее в работе "A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference" (433k примера)
Плюсы: поддерживает русский язык, можно конвентировать в GGUF.
Минусы: модели не очень хорошо понимают технический текст (особенно русский технический жаргон), нет версии в формате GGUF, и к числу фатальных недостатков могу отнести, что эти модели обучил не я ;)
Про выбор модели
Пришёл к выводу, что тему обучения подобных модей-эмбеддеров уже достаточно хорошо изучили и что можно без особых сложностей реализовать мою задумку.
В качестве базовой модели решил взять модель google-bert/bert-base-multilingual-uncased, потому что:
У этой крохи всего 168m параметров, что чуть больше чем у
paraphrase-multilingual-MiniLM-L12-v2, но меньше чем уparaphrase-multilingual-mpnet-base-v2;На вход она принимает 512 токенов, а на выходе возвращает 768-мерный эмбеддинг, столько же у
paraphrase-multilingual-mpnet-base-v2;Модель обучена на датасете wikipedia представляющем из себя Text Corpora, а там, сами понимаете, примеров текста больше, чем SNLI и Multi-Genre NLI вместе взятые;
Модель uncased, то есть обучение происходило на регистронезависимых текстах (сиречь всё переводилось в lowercase).
С моделью определились, теперь перейдём к вопросу выбора датасета...
Готовим датасет
Изначально я хотел собрать больше датасетов, но, собирая датасет по PHP, я понял, какой это трудоёмкий процесс, и решил уменьшить свои амбиции.
Итак, после поиска в интернете я нашёл только один подходящий датасет: OPUS PHP v1 на 2k примеров, содержащий пары текстов на русском и английском языках, по теме PHP.
Из указанного датасеста я использовал только английский корпус (так как русский корпус был очень низкого качества), далее задействовал инстанс LibreTranslate для перевода английских текстов на русский и очистил данные от аномалий и шума (сценарий dataset_php_build.ipynb). Затем вручную перевёл кривые места с помощью Google и Yandex Translate и экспортировал результат в CSV формат. Данные отсортировал и удалил дубликаты (сценарий dataset_php_undup.py) после чего осталось 1.6k примеров.
В финале попросил ChatGPT сгенерировать 100 примеров пар технического текста о PHP на русском и английском языках для сплита eval, а очищенные данные использовал для сплита train.
Результат выгрузил (сценарий dataset_php_publish.ipynb) на HuggingFace: evilfreelancer/opus-php-en-ru-cleaned .
Скрипт обучения на Domain Adaptation
Для создания эффективного эмбеддера, способного работать с техническими текстами о PHP на русском и английском языке, я решил провести обучение модели в два этапа, сначала выполнить Domain Adaptation, чтобы модель могла работать с техническими текстами на английском языке, а после этого обучить её на Parallel Corpora из русских и английских текстов.
Для Domain Adaptation я использовал метод Generative Pseudo Labeling (GPL) (arXiv:2112.07577), данный метод позволяет проводить обучение модели на основе неразмеченных данных, генерируя псевдометки и улучшая качество работы модели для специфических доменов.
Библиотека gpl имеет захардкоженный формат входного датасета и читает данные по определённым путям, поэтому пришлось слегка конвертировать тренировочный датасет и положить результат в директорию datasets (сценарий: dataset_php_convert.py).
Для адаптации модели bert-base-multilingual-uncased к домену английских текстов про PHP я использовал в качестве шаблона скрипт, предложенный авторами проекта GPL на их странице на GitHub, получился следующего вида код:
Полный скрипт тренировки train_domain.py можно найти в репозитории проекта на GitHub.
import gpl
model_name = 'bert-base-multilingual-uncased'
batch_size = 64
gpl_steps = 140000
output_dir = './output/enbeddrus_domain'
evaluation_output = f"{output_dir}_evaluation"
gpl.train(
path_to_generated_data=f"generated/embeddrus",
base_ckpt=model_name,
gpl_score_function="dot",
batch_size_gpl=batch_size,
gpl_steps=gpl_steps,
new_size=-1,
queries_per_passage=25,
output_dir=output_dir,
evaluation_data=f"./datasets",
evaluation_output=evaluation_output,
generator="BeIR/query-gen-msmarco-t5-base-v1",
retrievers=["msmarco-distilbert-base-v3", "msmarco-MiniLM-L-6-v3"],
retriever_score_functions=["cos_sim", "cos_sim"],
cross_encoder="cross-encoder/ms-marco-MiniLM-L-6-v2",
qgen_prefix="qgen",
do_evaluation=True,
)Процесс обучения включает в себя несколько этапов:
Используется генератор запросов, такой как BeIR/query-gen-msmarco-t5-base-v1, для создания синтетических запросов на основе текстов из корпуса;
С помощью ретриверов, таких как msmarco-distilbert-base-v3 и msmarco-MiniLM-L-6-v3, которые работают с косинусным сходством, извлекаются наиболее релевантные документы для сгенерированных запросов;
Кросс-энкодер, такой как cross-encoder/ms-marco-MiniLM-L-6-v2, используется для создания псевдометок, присваивая оценочные метки соответствия между запросами и документами;
Модель обучается с использованием MarginMSELoss, которая позволяет модели лучше адаптироваться к новому домену.
И так, наша модель обучена работать с новым доменом, поэтому переходим к следующему шагу.
Скрипт обучения на Parallel Corpora
Для обучения модели на параллельных корпусах я использовал метод обучения моделей на разных языках, описанный в примере на сайте Sentence Transformers. Этот метод позволяет обучать мультиязычные модели, используя параллельные тексты на разных языках (заготовка скрипта make_multilingual.py).
Для оценки качества модели я написал юпитер-блокнот, который загружает базовую и дообученную модель, прогоняет пары из eval сплита датасета evilfreelancer/opus-php-en-ru-cleaned и анализирует разницу между эмбеддингами, построенными для текстов на разных языках. Результаты визуализируются в виде графиков. Скрипт можно найти здесь.
На графике видно, что базовая модель bert-base-multilingual-uncased распределяет русские и английские тексты в изолированные кластеры точек, ну а наша задача сделать так, чтобы эти точки были расположены как можно ближе друг к другу.
Подобную задачу позволяет решать MSELoss, так как она минимизирует разницу между эмбеддингом, сгенерированным моделью-учителем (на английском языке) и эмбеддингом, сгенерированным моделью-учеником (на русском языке).
Используемые датасеты
Теперь пару слов про датасеты, решил остановиться на следующем наборе:
evilfreelancer/opus-php-en-ru-cleaned (1.6k) - ранее созданный датасет параллельных текстов на английском и русском языках;
Helsinki-NLP/opus_books (17.5k) - датасет OPUS параллельных текстов из книг.
Выбрал я их потому, что мои первые эксперименты с обучением модели на только PHP датасете показали, что у модели происходит overfitting в результате чего падала общее качество работы модели, поэтому самым логичным решением было добавить ещё один Parallel Corpora общего назначения.
Помимо этого в скрипт обучения я хотел сразу заложить возможность обучать на множестве разных датасетов (имеющих разные форматы данных), в результате чего получилась функция:
def read_datasets():
data = []
# Read cleaned OPUS PHP v1 en&ru dataset
docsphp_dataset = load_dataset("evilfreelancer/opus-php-en-ru-cleaned")
for item in docsphp_dataset['train']:
src_text = item["English"].strip()
trg_text = item["Russian"].strip()
if src_text and trg_text:
data.append((src_text, trg_text))
# Read OPUS Books v1 en&ru dataset
opus_dataset = load_dataset("Helsinki-NLP/opus_books", "en-ru")
for item in opus_dataset['train']:
src_text = item['translation']['en'].strip()
trg_text = item['translation']['ru'].strip()
if src_text and trg_text:
data.append((src_text, trg_text))
return dataВ дальнейшем планирую добавить в неё больше датасетов на разные технические темы, но на этапе прототипирования того что есть более чем достаточно.
Двигаемся дальше.
Обучение модели после Domain Adaptation
Полный скрипт тренировки train_parallel.py можно найти в репозитории проекта на GitHub, в качестве модели-учителя возьмём google-bert/bert-base-multilingual-uncased, а в качестве модели-ученика ту, что мы обучили ранее на шаге Domain Adaptation.
teacher_model_name = 'bert-base-multilingual-uncased'
student_model_name = './output/enbeddrus_domain'Обучение происходит в несколько этапов:
Сначала мы загружаем датасеты (функция read_datasets);
Далее выполняем их преобразование в нужный формат, после чего сохраняем на диске (функциия prepare_datasets)
Инициализируем модель-учитель и модель-ученик (тут)
Инициализируем MSELoss, передав ей на вход указатель на модель-ученика (тут)
Запускаем обучение модели-ученика
По завершению обучению давайте попробуем протестировать модель и понять стала ли на лучше извлекать эмбеддинги.
Как видно на графике эмбеддинги извлечённые из русских и английских текстов где-то наложились друг на друга, точность похожести поднялась с 0.83 до 0.94, при этом модель также хорошо разделяет фразы различающиеся по смыслу.
Веса обученной модели доступны тут: evilfreelancer/enbeddrus-v0.1-domain
Обучение модели только на Parallel Corpora
Посмотрел я на этот график и пришла в голову мысль, а что если попробовать обучить базовую модель сразу на Parallel Corpora, пропустив шаг с Domain Adaptation?
Правим скрипт тренировки, меняем модель-ученика, получается вот так:
teacher_model_name = 'bert-base-multilingual-uncased'
student_model_name = 'bert-base-multilingual-uncased'Опять запускаем тренировку и ждём некоторое время, по завершению прогоняем тесты и смотрим что получилось.
Как видно на графиках если обучать сразу на Parallel Corpora модель быстрее, так как не нужно выполнять Domain Adaptation, и лучше обучается извлекать эмбеддинги из параллельных текстов, ведь косинусное расстояние в таком случае между близкими по смыслу фразами на разных языках в среднем в районе 0.97, что выше чем у модели изначально обученной на домене текстов про PHP.
Веса обученной модели доступны тут: evilfreelancer/enbeddrus-v0.1
Отсюда можно сделать вывод, что дообучение мультиязыковой модели bert-base-multilingual-cased через Domain Adaptation с последующем обучением на Parallel Corpora не имеет особого смысла и проще сразу дообучать её на Parallel Corpora.
Финальный штрих
Осталось выполнить самую малость, для начала я хочу конвертировать модель в формат GGUF, чтобы можно было использовать обученные модели через llama.cpp, но на этом моменте сильно не будем заострять внимание, сошлюсь на мою публикацию "Как конвертировать модель BERT в формат GGUF?" в моём блоге и PR который я создал в проекте llama.cpp.
Но если кратко команды конвертации нужно выполнять и корня проекта llama.cpp и выглядят они следующим образом:
# Конвертируем модель Domain Adaptation + Parallel Corpora
python convert-hf-to-gguf.py \
../enbeddrus/output/enbeddrus-en-ru-2024-05-19_18-46-49 \
--outfile ../enbeddrus/models/enbeddrus-v0.1-domain-f16.gguf \
--outtype f16# Конвертируем модель Parallel Corpora
python convert-hf-to-gguf.py \
../enbeddrus/output/enbeddrus-en-ru-2024-05-20_11-30-48 \
--outfile ../enbeddrus/models/enbeddrus-v0.1-f16.gguf \
--outtype f16По её завершению в директории models появятся файлы: enbeddrus-v0.1-f16.gguf и enbeddrus-v0.1-domain-f16.gguf.
# Тегаем и публикуем Parallel Corpora
ollama create -f Modelfile.pc evilfreelancer/enbeddrus:latest
ollama push evilfreelancer/enbeddrus:latest
ollama create -f Modelfile.pc evilfreelancer/enbeddrus:v0.1
ollama push evilfreelancer/enbeddrus:v0.1
ollama create -f Modelfile.pc evilfreelancer/enbeddrus:v0.1-fp16
ollama push evilfreelancer/enbeddrus:v0.1-fp16Полученные модели я выгрузил на серверы Ollama следующим образом:
# Тегаем и публикуем Domain Adaptation + Parallel Corpora
ollama create -f Modelfile.dpc evilfreelancer/enbeddrus:v0.1-domain
ollama push evilfreelancer/enbeddrus:v0.1-domain
ollama create -f Modelfile.dpc evilfreelancer/enbeddrus:v0.1-domain-fp16
ollama push evilfreelancer/enbeddrus:v0.1-domain-fp16Выгруженные модели находятся тут и скачать их можно следующей командой:
ollama pull evilfreelancer/enbeddrusСодержимое Modelfile'ов можно найти в директории models проекта на GitHub.
Ссылки
Завершение
Благодаря работе над проектом enbeddrus были достигнуты следующие цели:
Удалось разобрался с тем как подобные модели устроены и как они работают, а так же с тем как их можно обучать;
Был собран датасет с Parallel Corpora тематических текстов о PHP на русском и английском;
Удалось разобраться с методами оценки моделей, а также с тем как эту оценку красиво визуализировать;
Была обучена модель, которая эффективно работает с текстами на двух языках и может быть использована в RAG-системе для поиска и анализа информации.
Полученные результаты подтверждают, что обучение мультиязычных эмбеддеров на основе параллельных корпусов является эффективным подходом для создания моделей, способных работать с текстами на разных языках.
Спасибо за внимание и за что дочитал публикацию до конца! Если у вас есть вопросы или вы хотите связаться со мной, ссылки на мои контакты в социальных сетях можно найти в моём профиле на Хабре.