train.py

import logging
import os
import wandb
import sys
from dataclasses import dataclass, field
from typing import Optional
from sklearn.metrics import f1_score, accuracy_score

import datasets
import numpy as np
from datasets import load_dataset

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from dataHelper import get_dataset
logger = logging.getLogger(__name__)



@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.

    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """

    dataset_name: Optional[str]=field(default=None, )
    max_seq_length: int=field(default=128, )

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str=field(default=None, )
    add_adapter: int=field(default=0, )


def main():

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    # parse a json file
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    
    wandb.init(
        project="NLPDL-Assignment2",

        config={
            "learning_rate": training_args.learning_rate,
            "epochs": training_args.num_train_epochs
        }
    )
    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # load the tokenizer from the arguments
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, )
    raw_datasets = get_dataset(data_args.dataset_name, tokenizer.sep_token)
    num_labels = max(raw_datasets['train']['labels']) + 1

    # Specify the number of labels and prepare the model
    config = AutoConfig.from_pretrained(model_args.model_name_or_path, num_labels=num_labels, )

    model = AutoModelForSequenceClassification.from_pretrained(model_args.model_name_or_path, config=config, )

    padding = "max_length"

    if data_args.max_seq_length > tokenizer.model_max_length:
        logger.warning(
            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
        )
    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    def preprocess_function(examples):
        result = tokenizer(examples['text'], padding=padding, max_length=max_seq_length, truncation=True)
        return result

    with training_args.main_process_first(desc="dataset map pre-processing"):
        raw_datasets = raw_datasets.map(
            preprocess_function,
            batched=True,
            desc="Running tokenizer on dataset",
        )

    train_dataset = raw_datasets["train"]
    eval_dataset = raw_datasets['test']

    # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
    # predictions and label_ids field) and has to return a dictionary string to float.
    def compute_metrics(p: EvalPrediction):
        preds = p.predictions
        preds = np.argmax(preds, axis=1)
        result = {
            "micro_f1": f1_score(p.label_ids, preds, average='micro'), 
            "macro_f1": f1_score(p.label_ids, preds, average='macro'), 
            "accuracy": accuracy_score(p.label_ids, preds),
        }
        return result

    # Data collator will default to DataCollatorWithPadding when the tokenizer is passed to Trainer, so we change it if
    # we already did the padding.
    data_collator = default_data_collator
    
    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # Training
    trainer.train()

    # Evaluation
    logger.info("*** Evaluate ***")
    metrics = trainer.evaluate(eval_dataset=eval_dataset)
    trainer.log_metrics("eval", metrics)
    trainer.save_metrics("eval", metrics)
    wandb.finish()

if __name__ == "__main__":
    main()