Осваиваем T5 (text-to-text transfer transformer). Fine-Tuning

Бывает, что при изучении материала по обучающей статье что-нибудь не работает, хотя коды копируются прямо из статьи.
В данном случае по обучающей статье был сделан Fine-Tuning модели T5 (text-to-text transfer transformer) по задаче машинного перевода, и в целом все получилось.

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

Работаем в Colab.
Переводим с английского языка на французский.

Install

Установка библиотек.
Внимание: обычная установка transformers выдает ошибку, понадобилась коррекция на вариант transformers[torch]

!pip install transformers[torch] datasets evaluate sacrebleu

Dataset

Скачиваем данные и разделяем на тренировочную и тестовую выборки.

from datasets import load_dataset
books = load_dataset("opus_books", "en-fr")
books = books["train"].train_test_split(test_size=0.2)

Посмотреть пример пары возможно так:

books["train"][0]

Preprocess

В Colab корректно сработали t5-small и t5-base.
На t5-large стандартным способом не хватило. памяти.

checkpoint = "t5-small"

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer

model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)

source_lang = "en"
target_lang = "fr"
prefix = "translate English to French: "

def preprocess_function(examples):
    inputs = [prefix + example[source_lang] for example in examples["translation"]]
    targets = [example[target_lang] for example in examples["translation"]]
    model_inputs = tokenizer(inputs, text_target=targets, max_length=128, truncation=True)
    return model_inputs

tokenized_books = books.map(preprocess_function, batched=True)

from transformers import DataCollatorForSeq2Seq
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=checkpoint)

Evaluate

Добавляем метрику BLEU.

import evaluate
metric = evaluate.load("sacrebleu")

import numpy as np

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [[label.strip()] for label in labels]
    return preds, labels

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

    result = metric.compute(predictions=decoded_preds, references=decoded_labels)
    result = {"bleu": result["score"]}

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
    result["gen_len"] = np.mean(prediction_lens)
    result = {k: round(v, 4) for k, v in result.items()}
    return result

Login

Для данного дообучения нужно авторизоваться, введя токен с правами "wtite".

from huggingface_hub import notebook_login
notebook_login()

Train

Загружаем выбранную предобученную модель

from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)

И формируем новую.
По сравнению со статьей добавлено "overwrite_output_dir=True", на случай перезаписи при сбоях.

new_model = 'my_t5_small_test'

training_args = Seq2SeqTrainingArguments(
    output_dir=new_model,
    overwrite_output_dir=True,
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=2,
    predict_with_generate=True,
    fp16=True,
    push_to_hub=True
)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_books["train"],
    eval_dataset=tokenized_books["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

Дообучение прошло успешно.

После завершения размещаем модель на HuggingFace.

trainer.push_to_hub()

Inference

text = "translate English to French: What is your name?"

from transformers import pipeline
translator = pipeline("translation", model=new_model)
translator(text)

>>> [{'translation_text': 'Quel est votre nom ?'}]J,hf

Обработчик сообщает, что хочет "translation_XX_to_YY" вместо "translation", но принципиальной разницы при коррекции не замечено.

Еще о собственных данных

Еще один датасет создавался для примера "вручную", не по статье.

texts = [
    {'en': 'The Wanderer', 'fr': 'Le grand Meaulnes'},
    {'en': 'Hello', 'fr': 'Bonjour'}
    ]

data_dict = {'id': [ key for key in range(len(texts)) ], 'translation': texts}
my_dataset = Dataset.from_dict(data_dict)
dataset_dict = DatasetDict({"train": my_dataset})

Так тоже все сработало.
Из этого фрагмента понятно, как составить датасет независимо от того, в каком виде пары находятся изначально. В любом случае возможно сформировать массив texts с помощью циклов и некоторых преобразований.

Вопросы, оставшиеся неясными

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

Примечания

Если Вы обнаружили в статье неточность, или считаете полезным что-либо добавить - пожалуйста, сообщите в комментариях.