zeroshot_training_script.py

#!/usr/bin/env python
# coding: utf-8

# # Creating a Zero-Shot classifier based on BETO
# 
# This notebook/script fine-tunes a BETO (spanish bert, 'dccuchile/bert-base-spanish-wwm-cased') model on the spanish XNLI dataset.
# The fine-tuned model can then be fed to a Huggingface ZeroShot pipeline to obtain a ZeroShot classifier.

# In[ ]:


from datasets import load_dataset, Dataset, load_metric, load_from_disk
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import Trainer, TrainingArguments
import torch
from pathlib import Path
# from ray import tune
# from ray.tune.suggest.hyperopt import HyperOptSearch
# from ray.tune.schedulers import ASHAScheduler


# # Prepare the datasets

# In[ ]:


xnli_es = load_dataset("xnli", "es")


# In[ ]:


xnli_es


# >joeddav
# >Aug '20
# >
# >@rsk97 In addition, just make sure the model used is trained on an NLI task and that the **last output label corresponds to entailment** while the **first output label corresponds to contradiction**.
# 
# => We change the original `label` and use the `labels` column, which is required by a `AutoModelForSequenceClassification`

# In[ ]:


# see markdown above
def switch_label_id(row):
    if row["label"] == 0:
        return {"labels": 2}
    elif row["label"] == 2:
        return {"labels": 0}
    else:
        return {"labels": 1}

for split in xnli_es:
    xnli_es[split] = xnli_es[split].map(switch_label_id)


# ## Tokenize data

# In[ ]:


tokenizer = AutoTokenizer.from_pretrained("dccuchile/bert-base-spanish-wwm-cased")


# In a first attempt i padded all data to the maximum length of the dataset (379). However, the traninig takes substanially longer with all the paddings, it's better to pass in the tokenizer to the `Trainer` and let the `Trainer` do the padding on a batch level.

# In[ ]:


# Figured out max length of the dataset manually
# max_length = 379
def tokenize(row):
    return tokenizer(row["premise"], row["hypothesis"], truncation=True, max_length=512)  #, padding="max_length", max_length=max_length)


# In[ ]:


data = {}
for split in xnli_es:
    data[split] = xnli_es[split].map(
        tokenize, 
        remove_columns=["hypothesis", "premise", "label"], 
        batched=True, 
        batch_size=128
    )


# In[ ]:


train_path = str(Path("./train_ds").absolute())
valid_path = str(Path("./valid_ds").absolute())

data["train"].save_to_disk(train_path)
data["validation"].save_to_disk(valid_path)


# In[ ]:


# We can use `datasets.Dataset`s directly

# class XnliDataset(torch.utils.data.Dataset):
#     def __init__(self, data):
#         self.data = data

#     def __getitem__(self, idx):
#         item = {key: torch.tensor(val) for key, val in self.data[idx].items()}
#         return item

#     def __len__(self):
#         return len(self.data)


# In[ ]:


def trainable(config):
    metric = load_metric("xnli", "es")

    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        predictions = predictions.argmax(axis=-1)
        return metric.compute(predictions=predictions, references=labels)
    
    model = AutoModelForSequenceClassification.from_pretrained("dccuchile/bert-base-spanish-wwm-cased", num_labels=3)

    training_args = TrainingArguments(
        output_dir='./results',          # output directory
        do_train=True,
        do_eval=True,
        evaluation_strategy="steps",
        eval_steps=500,
        load_best_model_at_end=True,
        metric_for_best_model="eval_accuracy",
        num_train_epochs=config["epochs"],              # total number of training epochs
        per_device_train_batch_size=config["batch_size"],  # batch size per device during training
        per_device_eval_batch_size=config["batch_size_eval"],   # batch size for evaluation
        warmup_steps=config["warmup_steps"],  # 500
        weight_decay=config["weight_decay"],  # 0.001               # strength of weight decay
        learning_rate=config["learning_rate"],  # 5e-05
        logging_dir='./logs',            # directory for storing logs
        logging_steps=250,
        #save_steps=500,  # ignored when using load_best_model_at_end
        save_total_limit=10,
        no_cuda=False,
        disable_tqdm=True,
    )
    
#     train_dataset = XnliDataset(load_from_disk(config["train_path"]))
#     valid_dataset = XnliDataset(load_from_disk(config["valid_path"]))
    train_dataset = load_from_disk(config["train_path"])
    valid_dataset = load_from_disk(config["valid_path"])

    
    trainer = Trainer(
        model,
        tokenizer=tokenizer,
        args=training_args,                  # training arguments, defined above
        train_dataset=train_dataset,         # training dataset
        eval_dataset=valid_dataset,          # evaluation dataset
        compute_metrics=compute_metrics,
    )
    
    trainer.train()


# In[ ]:


trainable(
    {
        "train_path": train_path,
        "valid_path": valid_path,
        "batch_size": 16,
        "batch_size_eval": 64,
        "warmup_steps": 500,
        "weight_decay": 0.001,
        "learning_rate": 5e-5,
        "epochs": 3,
    }
)


# # HPO

# In[ ]:


# config = {
#     "train_path": train_path,
#     "valid_path": valid_path,
#     "warmup_steps": tune.randint(0, 500),
#     "weight_decay": tune.loguniform(0.00001, 0.1),
#     "learning_rate": tune.loguniform(5e-6, 5e-4),
#     "epochs": tune.choice([2, 3, 4])
# }


# # In[ ]:


# analysis = tune.run(
#     trainable,
#     config=config,
#     metric="eval_acc",
#     mode="max",
#     #search_alg=HyperOptSearch(),
#     #scheduler=ASHAScheduler(),
#     num_samples=1,
# )


# # In[ ]:


# def model_init():
#     return AutoModelForSequenceClassification.from_pretrained("dccuchile/bert-base-spanish-wwm-cased", num_labels=3)

# trainer = Trainer(
#     args=training_args,                  # training arguments, defined above
#     train_dataset=train_dataset,         # training dataset
#     eval_dataset=valid_dataset,          # evaluation dataset
#     model_init=model_init,
#     compute_metrics=compute_metrics,
# )


# # In[ ]:


# best_trial = trainer.hyperparameter_search(
#     direction="maximize",
#     backend="ray",
#     n_trials=2,
#     # Choose among many libraries:
#     # https://docs.ray.io/en/latest/tune/api_docs/suggestion.html
#     search_alg=HyperOptSearch(mode="max", metric="accuracy"),
#     # Choose among schedulers:
#     # https://docs.ray.io/en/latest/tune/api_docs/schedulers.html
#     scheduler=ASHAScheduler(mode="max", metric="accuracy"),
#     local_dir="tune_runs",
# )