Небольшая языковая модель для генерации анекдотов

Учебный проект по реализации Byte-level BPE токенизатора и Transformer с нуля. Модель обучена предсказывать следующий токен на датасете русских анекдотов. В этом репозитории размещена версия mini после основного обучения и дообучения.

Что реализовано

  • Обучение Byte-level BPE, кодирование и декодирование текста.
  • Transformer с RMSNorm, позиционными смещениями ALiBi, Grouped-Query Attention и SwiGLU.
  • Обучение с маскированием паддингов в loss, warmup и линейным уменьшением learning rate.
  • Генерация с temperature и top-k sampling.

Архитектура mini

Параметр Значение
Число слоёв 6
Размер скрытого представления 384
Число голов запросов 6
Число голов ключей и значений 3
Промежуточный размер SwiGLU 1024
Размер словаря 1024, включая [EOS]
Максимальный контекст 128 токенов
Dropout 0,1

Данные

Использован датасет IgorVolochay/russian_jokes, загруженный из dataset.txt через загрузчик text.

Данные случайно разделены в отношении 90/10 с seed=0xC0FFEE: 135 497 текстов для обучения и 15 056 для оценки. Токенизатор обучался только на тренировочной части. Средняя длина текста на отложенной части — 70,42 токена. Последовательности длиннее 128 токенов при обучении обрезались.

Обучение

Обучение выполнялось в Google Colab на CUDA с batch size 16, оптимизатором AdamW, weight decay 0,01 и ограничением нормы градиента до 1,0. Функция потерь — cross-entropy для предсказания следующего токена. Паддинги исключены из расчёта, настоящий EOS учитывается.

Этап Число шагов Максимальный learning rate
Обучение mini с нуля 10 000 0,0003
Дообучение тех же весов 20 000 0,0001

В начале каждого этапа learning rate линейно увеличивался в течение 10% шагов, затем линейно уменьшался до нуля. При дообучении оптимизатор и scheduler создавались заново, веса модели сохранялись. Всего mini прошла 30 000 шагов, или примерно 3,54 прохода по тренировочным данным.

Результаты

Модель и этап Итоговый validation loss
nano, 10 000 шагов 3,616
mini, 10 000 шагов 2,963
mini, ещё 20 000 шагов 2,783

Значения приведены по сохранённым результатам ноутбука. Валидация использует среднее значение loss по батчам. Во время дообучения validation loss снижался, но улучшение к концу замедлялось. Отложенная часть использовалась для сравнения экспериментов; отдельная финальная тестовая выборка не выделялась.

Сравнение генерации

На пяти одинаковых началах сравнены настройки temperature=0.8, top_k=10 и temperature=0.6, top_k=5. Максимум продолжения — 100 новых токенов. Для каждого начала перед обеими генерациями задавался одинаковый seed.

Начало: «Мужик спрашивает».

При temperature=0.8, top_k=10:

Мужик спрашивает у своего друга:- Слушай, а почему у тебя пахнет с собой?- По дереву, а ты - кусочек, а у тебя - ногу. А с ним - ничего.

При temperature=0.6, top_k=5:

Мужик спрашивает у жены: - Дорогой, что ты делаешь? - Я не знаю, я хочу, чтобы тебя не доверяли.

При более низкой температуре отдельные фразы выглядели аккуратнее, но в других примерах усиливались повторы вроде «вчера вчера» и «вечером вечером». Устойчивого преимущества одной настройки не обнаружено. Это небольшая ручная проверка, а не статистическая оценка качества.

Как запустить

Нужны определения классов ByteLevelBPETokenizer и TransformerForCausalLM, а также их зависимостей из учебного ноутбука. Архитектура реализована вручную; пример использует эти классы.

import torch

repo_id = "MayaLager/llm-course-hw1-mini"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = ByteLevelBPETokenizer.from_pretrained(repo_id)
model = TransformerForCausalLM.from_pretrained(repo_id).to(device)
model.eval()

input_ids = torch.tensor(
    tokenizer.encode("Заходит в бар", add_eos_token=False),
    device=device,
)[None, :]

output = model.generate(
    input_ids,
    max_new_tokens=100,
    eos_token_id=tokenizer.eos_token_id,
    temperature=0.8,
    do_sample=True,
    top_k=10,
)

print(tokenizer.decode(output[0].tolist()))

Выводы и ограничения

Переход от nano к mini и дальнейшее дообучение улучшили validation loss. Модель воспроизводит форму диалогов и характерные начала анекдотов, но пока не строит устойчивый сюжет и осмысленную шутку. В генерациях встречаются нарушения согласования, смена персонажей, повторы и незаконченные слова при достижении лимита токенов.

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

Downloads last month
115
Safetensors
Model size
10.6M params
Tensor type
F32
·
BOOL
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train MayaLager/llm-course-hw1-mini