src/train.py

from preprocess import PreprocessingDatasetArguments
from shared import CustomTokens, load_datasets, CustomTrainingArguments, get_last_checkpoint, train_from_checkpoint
from model import ModelArguments
import transformers
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional
from datasets import utils as d_utils
from transformers import (
    DataCollatorForSeq2Seq,
    HfArgumentParser,
    Seq2SeqTrainer,
)

from transformers.utils import check_min_version
from transformers.utils.versions import require_version

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version('4.13.0.dev0')
require_version('datasets>=1.8.0',
                'To fix: pip install -r requirements.txt')

os.environ['WANDB_DISABLED'] = 'true'

logging.basicConfig()
logger = logging.getLogger(__name__)

# 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)],
)



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

    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={'help': 'The number of processes to use for the preprocessing.'},
    )

    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            'help': 'For debugging purposes or quicker training, truncate the number of training examples to this value if set.'
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            'help': 'For debugging purposes or quicker training, truncate the number of evaluation examples to this value if set.'
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            'help': 'For debugging purposes or quicker training, truncate the number of prediction examples to this value if set.'
        },
    )


def main():

    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    hf_parser = HfArgumentParser((
        ModelArguments,
        PreprocessingDatasetArguments,
        DataTrainingArguments,
        CustomTrainingArguments
    ))
    model_args, dataset_args, data_training_args, training_args = hf_parser.parse_args_into_dataclasses()

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    d_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()

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

    # 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}')

    # FP16 https://github.com/huggingface/transformers/issues/9295

    # Works:
    # https://huggingface.co/docs/transformers/model_doc/t5v1.1
    # google/t5-v1_1-small
    # google/t5-v1_1-base
    # google/t5-v1_1-large
    # google/t5-v1_1-xl
    # google/t5-v1_1-xxl

    # https://huggingface.co/docs/transformers/model_doc/t5
    # t5-small
    # t5-base
    # t5-large
    # t5-3b
    # t5-11b

    # allenai/led-base-16384 - https://github.com/huggingface/transformers/issues/9810

    # Further work:
    # Multilingual- https://huggingface.co/docs/transformers/model_doc/mt5

    # In distributed training, the load_dataset function guarantees that only one local process can concurrently
    # download the dataset.
    raw_datasets = load_datasets(dataset_args)
    # , cache_dir=model_args.cache_dir

    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.


    # Detecting last checkpoint.
    last_checkpoint = get_last_checkpoint(training_args)

    from model import get_model_tokenizer
    model, tokenizer = get_model_tokenizer(model_args, training_args)

    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    column_names = raw_datasets['train'].column_names

    prefix = CustomTokens.EXTRACT_SEGMENTS_PREFIX.value

    PAD_TOKEN_REPLACE_ID = -100

    # https://github.com/huggingface/transformers/issues/5204
    def preprocess_function(examples):
        inputs = examples['text']
        targets = examples['extracted']
        inputs = [prefix + inp for inp in inputs]
        model_inputs = tokenizer(inputs, truncation=True)

        # Setup the tokenizer for targets
        with tokenizer.as_target_tokenizer():
            labels = tokenizer(targets, truncation=True)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100
        # when we want to ignore padding in the loss.

        model_inputs['labels'] = [
            [(l if l != tokenizer.pad_token_id else PAD_TOKEN_REPLACE_ID)
                for l in label]
            for label in labels['input_ids']
        ]

        return model_inputs

    def prepare_dataset(dataset, desc):
        return dataset.map(
            preprocess_function,
            batched=True,
            num_proc=data_training_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not dataset_args.overwrite_cache,
            desc=desc,  # tokenizing train dataset
        )
    # train_dataset # TODO shuffle?

    # if training_args.do_train:
    if 'train' not in raw_datasets:  # TODO do checks above?
        raise ValueError('Train dataset missing')
    train_dataset = raw_datasets['train']
    if data_training_args.max_train_samples is not None:
        train_dataset = train_dataset.select(
            range(data_training_args.max_train_samples))
    with training_args.main_process_first(desc='train dataset map pre-processing'):
        train_dataset = prepare_dataset(
            train_dataset, desc='Running tokenizer on train dataset')

    if 'validation' not in raw_datasets:
        raise ValueError('Validation dataset missing')
    eval_dataset = raw_datasets['validation']
    if data_training_args.max_eval_samples is not None:
        eval_dataset = eval_dataset.select(
            range(data_training_args.max_eval_samples))
    with training_args.main_process_first(desc='validation dataset map pre-processing'):
        eval_dataset = prepare_dataset(
            eval_dataset, desc='Running tokenizer on validation dataset')

    if 'test' not in raw_datasets:
        raise ValueError('Test dataset missing')
    predict_dataset = raw_datasets['test']
    if data_training_args.max_predict_samples is not None:
        predict_dataset = predict_dataset.select(
            range(data_training_args.max_predict_samples))
    with training_args.main_process_first(desc='prediction dataset map pre-processing'):
        predict_dataset = prepare_dataset(
            predict_dataset, desc='Running tokenizer on prediction dataset')

    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=PAD_TOKEN_REPLACE_ID,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )

    # Done processing datasets

    # Initialize our Trainer
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # Training
    train_result = train_from_checkpoint(trainer, last_checkpoint, training_args)

    metrics = train_result.metrics
    max_train_samples = data_training_args.max_train_samples or len(
        train_dataset)
    metrics['train_samples'] = min(max_train_samples, len(train_dataset))

    trainer.log_metrics('train', metrics)
    trainer.save_metrics('train', metrics)
    trainer.save_state()

    kwargs = {'finetuned_from': model_args.model_name_or_path,
                'tasks': 'summarization'}

    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:
        trainer.create_model_card(**kwargs)


if __name__ == '__main__':
    main()