

Before you begin, make sure you have all the necessary libraries installed:

pip install transformers datasets evaluate sacrebleu
We encourage you to login to your Hugging Face account so you can upload and share your model with the community. When prompted, enter your token to login:

from huggingface_hub import notebook_login
notebook_login()

Load OPUS Books dataset
Start by loading the English-French subset of the OPUS Books dataset from the 🤗 Datasets library:

from datasets import load_dataset
books = load_dataset("opus_books", "en-fr")

Split the dataset into a train and test set with the [~datasets.Dataset.train_test_split] method:

books = books["train"].train_test_split(test_size=0.2)

Then take a look at an example:

books["train"][0]
{'id': '90560',
 'translation': {'en': 'But this lofty plateau measured only a few fathoms, and soon we reentered Our Element.',
  'fr': 'Mais ce plateau élevé ne mesurait que quelques toises, et bientôt nous fûmes rentrés dans notre élément.'}}

translation: an English and French translation of the text.
Preprocess

The next step is to load a T5 tokenizer to process the English-French language pairs:

from transformers import AutoTokenizer
checkpoint = "t5-small"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

The preprocessing function you want to create needs to:

Prefix the input with a prompt so T5 knows this is a translation task. Some models capable of multiple NLP tasks require prompting for specific tasks.
Tokenize the input (English) and target (French) separately because you can't tokenize French text with a tokenizer pretrained on an English vocabulary.
Truncate sequences to be no longer than the maximum length set by the max_length parameter.

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

To apply the preprocessing function over the entire dataset, use 🤗 Datasets [~datasets.Dataset.map] method. You can speed up the map function by setting batched=True to process multiple elements of the dataset at once:

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

Now create a batch of examples using [DataCollatorForSeq2Seq]. It's more efficient to dynamically pad the sentences to the longest length in a batch during collation, instead of padding the whole dataset to the maximum length.

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

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

Evaluate
Including a metric during training is often helpful for evaluating your model's performance. You can quickly load a evaluation method with the 🤗 Evaluate library. For this task, load the SacreBLEU metric (see the 🤗 Evaluate quick tour to learn more about how to load and compute a metric):

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

Then create a function that passes your predictions and labels to [~evaluate.EvaluationModule.compute] to calculate the SacreBLEU score:

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

Your compute_metrics function is ready to go now, and you'll return to it when you setup your training.
Train

If you aren't familiar with finetuning a model with the [Trainer], take a look at the basic tutorial here!

You're ready to start training your model now! Load T5 with [AutoModelForSeq2SeqLM]:

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

At this point, only three steps remain:

Define your training hyperparameters in [Seq2SeqTrainingArguments]. The only required parameter is output_dir which specifies where to save your model. You'll push this model to the Hub by setting push_to_hub=True (you need to be signed in to Hugging Face to upload your model). At the end of each epoch, the [Trainer] will evaluate the SacreBLEU metric and save the training checkpoint.
Pass the training arguments to [Seq2SeqTrainer] along with the model, dataset, tokenizer, data collator, and compute_metrics function.
Call [~Trainer.train] to finetune your model.

training_args = Seq2SeqTrainingArguments(
     output_dir="my_awesome_opus_books_model",
     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()
`

Once training is completed, share your model to the Hub with the [~transformers.Trainer.push_to_hub] method so everyone can use your model:

trainer.push_to_hub()

If you aren't familiar with finetuning a model with Keras, take a look at the basic tutorial here!

To finetune a model in TensorFlow, start by setting up an optimizer function, learning rate schedule, and some training hyperparameters:

from transformers import AdamWeightDecay
optimizer = AdamWeightDecay(learning_rate=2e-5, weight_decay_rate=0.01)

Then you can load T5 with [TFAutoModelForSeq2SeqLM]:

from transformers import TFAutoModelForSeq2SeqLM
model = TFAutoModelForSeq2SeqLM.from_pretrained(checkpoint)

Convert your datasets to the tf.data.Dataset format with [~transformers.TFPreTrainedModel.prepare_tf_dataset]:

tf_train_set = model.prepare_tf_dataset(
     tokenized_books["train"],
     shuffle=True,
     batch_size=16,
     collate_fn=data_collator,
 )
tf_test_set = model.prepare_tf_dataset(
     tokenized_books["test"],
     shuffle=False,
     batch_size=16,
     collate_fn=data_collator,
 )

Configure the model for training with compile. Note that Transformers models all have a default task-relevant loss function, so you don't need to specify one unless you want to:

import tensorflow as tf
model.compile(optimizer=optimizer)  # No loss argument!

The last two things to setup before you start training is to compute the SacreBLEU metric from the predictions, and provide a way to push your model to the Hub. Both are done by using Keras callbacks.
Pass your compute_metrics function to [~transformers.KerasMetricCallback]:

from transformers.keras_callbacks import KerasMetricCallback
metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_validation_set)

Specify where to push your model and tokenizer in the [~transformers.PushToHubCallback]:

from transformers.keras_callbacks import PushToHubCallback
push_to_hub_callback = PushToHubCallback(
     output_dir="my_awesome_opus_books_model",
     tokenizer=tokenizer,
 )

Then bundle your callbacks together:

callbacks = [metric_callback, push_to_hub_callback]

Finally, you're ready to start training your model! Call fit with your training and validation datasets, the number of epochs, and your callbacks to finetune the model:

model.fit(x=tf_train_set, validation_data=tf_test_set, epochs=3, callbacks=callbacks)

Once training is completed, your model is automatically uploaded to the Hub so everyone can use it!

For a more in-depth example of how to finetune a model for translation, take a look at the corresponding
PyTorch notebook
or TensorFlow notebook.

Inference
Great, now that you've finetuned a model, you can use it for inference!
Come up with some text you'd like to translate to another language. For T5, you need to prefix your input depending on the task you're working on. For translation from English to French, you should prefix your input as shown below:

text = "translate English to French: Legumes share resources with nitrogen-fixing bacteria."

The simplest way to try out your finetuned model for inference is to use it in a [pipeline]. Instantiate a pipeline for translation with your model, and pass your text to it:

from transformers import pipeline
translator = pipeline("translation", model="my_awesome_opus_books_model")
translator(text)
[{'translation_text': 'Legumes partagent des ressources avec des bactéries azotantes.'}]

You can also manually replicate the results of the pipeline if you'd like:

Tokenize the text and return the input_ids as PyTorch tensors:

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("my_awesome_opus_books_model")
inputs = tokenizer(text, return_tensors="pt").input_ids

Use the [~transformers.generation_utils.GenerationMixin.generate] method to create the translation. For more details about the different text generation strategies and parameters for controlling generation, check out the Text Generation API.

from transformers import AutoModelForSeq2SeqLM
model = AutoModelForSeq2SeqLM.from_pretrained("my_awesome_opus_books_model")
outputs = model.generate(inputs, max_new_tokens=40, do_sample=True, top_k=30, top_p=0.95)

Decode the generated token ids back into text:

tokenizer.decode(outputs[0], skip_special_tokens=True)
'Les lignées partagent des ressources avec des bactéries enfixant l'azote.'
``
</pt>
<tf>
Tokenize the text and return theinput_ids` as TensorFlow tensors:

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("my_awesome_opus_books_model")
inputs = tokenizer(text, return_tensors="tf").input_ids

Use the [~transformers.generation_tf_utils.TFGenerationMixin.generate] method to create the translation. For more details about the different text generation strategies and parameters for controlling generation, check out the Text Generation API.

from transformers import TFAutoModelForSeq2SeqLM
model = TFAutoModelForSeq2SeqLM.from_pretrained("my_awesome_opus_books_model")
outputs = model.generate(inputs, max_new_tokens=40, do_sample=True, top_k=30, top_p=0.95)

Decode the generated token ids back into text:

tokenizer.decode(outputs[0], skip_special_tokens=True)
'Les lugumes partagent les ressources avec des bactéries fixatrices d'azote.'

