Enbeddrus — обучение независящей от языка эмбеддинг-модели02.06.2024 18:00

Приветствую, хабровчане!

Сегодня я хочу рассказать вам историю о том, как я обучил простую и компактную независящую от языка (language agnostic) модель-эмбеддер, которая умеет работать с техническими текстами о PHP и способна извлекать схожие эмбеддинги для параллельных текстов на английском и русском языках.

Вот что предложила модель DALL-E в качестве иллюстрации к данной публикации

Вот что предложила модель DALL-E в качестве иллюстрации к данной публикации

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

Для достижения этой цели как-раз и необходима независимая от языка модель-эмбеддер, которая будет одинаково хорошо работать с техническими текстами на русском и английском языках.

Кратко о том, как работает модель Enbeddrus

Кратко о том, как работает модель Enbeddrus

Ещё одним важным аспектом было то, чтобы модель потребляла как можно меньше ресурсов и, если возможно, чтобы её можно было преобразовать в формат 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-мерный эмбеддинг.

Обе эти модели обучены на комбинации из датасетов:

Плюсы: поддерживает русский язык, можно конвентировать в 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,
)

Процесс обучения включает в себя несколько этапов:

  1. Используется генератор запросов, такой как BeIR/query-gen-msmarco-t5-base-v1, для создания синтетических запросов на основе текстов из корпуса;

  2. С помощью ретриверов, таких как msmarco-distilbert-base-v3 и msmarco-MiniLM-L-6-v3, которые работают с косинусным сходством, извлекаются наиболее релевантные документы для сгенерированных запросов;

  3. Кросс-энкодер, такой как cross-encoder/ms-marco-MiniLM-L-6-v2, используется для создания псевдометок, присваивая оценочные метки соответствия между запросами и документами;

  4. Модель обучается с использованием MarginMSELoss, которая позволяет модели лучше адаптироваться к новому домену.

И так, наша модель обучена работать с новым доменом, поэтому переходим к следующему шагу.

Скрипт обучения на Parallel Corpora

Для обучения модели на параллельных корпусах я использовал метод обучения моделей на разных языках, описанный в примере на сайте Sentence Transformers. Этот метод позволяет обучать мультиязычные модели, используя параллельные тексты на разных языках (заготовка скрипта make_multilingual.py).

Для оценки качества модели я написал юпитер-блокнот, который загружает базовую и дообученную модель, прогоняет пары из eval сплита датасета evilfreelancer/opus-php-en-ru-cleaned и анализирует разницу между эмбеддингами, построенными для текстов на разных языках. Результаты визуализируются в виде графиков. Скрипт можно найти здесь.

TSNE распределение эмбеддингов базовой мультиязыковой модели

TSNE распределение эмбеддингов базовой мультиязыковой модели

На графике видно, что базовая модель bert-base-multilingual-uncased распределяет русские и английские тексты в изолированные кластеры точек, ну, а наша задача сделать так, чтобы эти точки были расположены как можно ближе друг к другу.

Подобную задачу позволяет решать MSELoss, так как она минимизирует разницу между эмбеддингом, сгенерированным моделью-учителем (на английском языке) и эмбеддингом, сгенерированным моделью-учеником (на русском языке).

Используемые датасеты

Теперь пару слов про датасеты, решил остановиться на следующем наборе:

Выбрал я их потому, что мои первые эксперименты с обучением модели на только 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'

Обучение происходит в несколько этапов:

  1. Сначала мы загружаем датасеты (функция read_datasets);

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

  3. Инициализируем модель-учитель и модель-ученик (тут)

  4. Инициализируем MSELoss, передав ей на вход указатель на модель-ученика (тут)

  5. Запускаем обучение модели-ученика

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

Графики TSNE и косинусного расстояния модели Domain Adaptation + Parallel Corpora

Графики TSNE и косинусного расстояния модели Domain Adaptation + Parallel Corpora

Как видно на графике эмбеддинги извлечённые из русских и английских текстов где-то наложились друг на друга, точность похожести поднялась с 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'

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

Графики TSNE и косинусного расстояния модели Parallel Corpora

Графики TSNE и косинусного расстояния модели Parallel Corpora

Как видно на графиках если обучать сразу на 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 были достигнуты следующие цели:

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

  2. Был собран датасет с Parallel Corpora тематических текстов о PHP на русском и английском;

  3. Удалось разобраться с методами оценки моделей, а также с тем как эту оценку красиво визуализировать;

  4. Была обучена модель, которая эффективно работает с текстами на двух языках и может быть использована в RAG-системе для поиска и анализа информации.

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

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

© Habrahabr.ru