Finetuning a Hugging Face model

In this guide, we finetune a Hugging Face model on the tweet_eval emoji task. The transformers library contains all code to do the training for us, we only need to define a torch Dataset that contains our data and handles tokenization.

# Hyperparameters
model_name = "distilroberta-base"
batch_size = 32
learning_rate = 1e-3

class TransformerDataset(torch.utils.data.Dataset):
    def __init__(self, data: pd.DataFrame, tokenizer: transformers.PreTrainedTokenizerBase):
        self.data = data
        self.label2idx, self.idx2label = self.get_label_map()
        self.num_labels = len(self.label2idx)
        self.tokenizer = tokenizer
        
    def tokenize(self, message, label=None):
        tokenized = self.tokenizer(message, padding="max_length", truncation=True)
        if label:
            tokenized["label"] = self.label2idx[label]
        return tokenized

    def get_label_map(self):
        unique_labels = data["annotation.labelValue"].unique()
        return {l: i for i, l in enumerate(unique_labels)}, {i: l for i, l in enumerate(unique_labels)}
        
    def __len__(self):
        return len(self.data)
        
    def __getitem__(self, idx):
        # Get the row from self.data, and skip the first column (id).
        return self.tokenize(*self.data.iloc[idx][1:])

tokenizer = AutoTokenizer.from_pretrained(model_name)
dataset = TransformerDataset(data, tokenizer)

Downloading:   0%|          | 0.00/480 [00:00<?, ?B/s]



Downloading:   0%|          | 0.00/878k [00:00<?, ?B/s]



Downloading:   0%|          | 0.00/446k [00:00<?, ?B/s]



Downloading:   0%|          | 0.00/1.29M [00:00<?, ?B/s]

The 🤗 Transformers library provides all the code we need for training a RoBERTa model. We will use their Trainer class, which handles all training, monitoring and integration with Weights & Biases for us. The 🤗 Transformers documentation has a detailed tutorial on fine-tuning models, which can be found here.

# Load model
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=dataset.num_labels,
    id2label=dataset.idx2label
)

# To increase training speed, we will freeze all layers except the classifier head.
for param in model.base_model.parameters():
    param.requires_grad = False
training_args = transformers.TrainingArguments(
    "twitter-emoji-trainer",
    learning_rate=learning_rate,
    per_device_train_batch_size=batch_size,
    logging_steps=1,
    optim="adamw_torch"
)

trainer = transformers.Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset
)
logging.set_verbosity(40)
trainer.train()

Downloading:   0%|          | 0.00/316M [00:00<?, ?B/s]


{'loss': 0.7297, 'learning_rate': 0.0006666666666666666, 'epoch': 1.0}
{'loss': 0.732, 'learning_rate': 0.0003333333333333333, 'epoch': 2.0}
{'loss': 0.7407, 'learning_rate': 0.0, 'epoch': 3.0}
{'train_runtime': 0.6766, 'train_samples_per_second': 17.736, 'train_steps_per_second': 4.434, 'train_loss': 0.7341360052426656, 'epoch': 3.0}





TrainOutput(global_step=3, training_loss=0.7341360052426656, metrics={'train_runtime': 0.6766, 'train_samples_per_second': 17.736, 'train_steps_per_second': 4.434, 'train_loss': 0.7341360052426656, 'epoch': 3.0})

Uploading your model

We trained our model. The last step is to upload the model to the package registry of our repo at gitlab.memri.io, so we can use it in a plugin. To do this, you are required to create a personal project on gitlab.memri.io for your plugin. For this example, we created sentiment-plugin to upload our model to. When running the command below, you will be asked to provide a personal_acces_token (just follow the link), which will be used to upload the model to gitlab on your behalf.

write_model_to_package_registry(model, project_name="sentiment-plugin")

        The first time you are uploading a model you need to create an access_token
        at https://gitlab.memri.io/-/profile/personal_access_tokens?name=Model+Access+token&scopes=api
        Click at the blue button with 'Create personal access token'"
        
Then copy your personal access token from 'Your new personal access token', and paste here: ··········
writing config.json to package registry of sentiment-plugin with project id 175
uploading /tmp/config.json

Succesfully uploaded /tmp/config.json
writing pytorch_model.bin to package registry of sentiment-plugin with project id 175
uploading /tmp/pytorch_model.bin

Succesfully uploaded /tmp/pytorch_model.bin

That’s it, we created our model and made it accesible via the package registry. Check out the next tutorial to see how you can use this model from within your data app.