finetune_text2text-Copy1.py

# +
import torch
from transformers import ByT5Tokenizer, T5Tokenizer, T5Config
import datasets

tokenizer = ByT5Tokenizer.from_pretrained('google/byt5-small',use_fast = False)
#tokenizer = T5Tokenizer.from_pretrained('t5-small',use_fast = False)

# process the examples in input and target text format and the eos token at the end 
def add_eos_to_examples(example):

    example['input_text'] = 'question: %s  context: %s </s>' % (example['asr_question_parsed'], example['asr_content_parsed'])
    example['target_text'] = '%s </s>' % example['asr_answer_parsed']
    
    return example

# tokenize the examples
def convert_to_features(example_batch):
    input_encodings = tokenizer.batch_encode_plus(example_batch['input_text'], padding='max_length', truncation=True, max_length=1024)
    target_encodings = tokenizer.batch_encode_plus(example_batch['target_text'], padding='max_length', truncation=True, max_length=128)
    #print(input_encodings.keys())
    #print(target_encodings.keys())
    #s()
    encodings = {
        'input_ids': input_encodings['input_ids'], 
        'attention_mask': input_encodings['attention_mask'],
        'target_ids': target_encodings['input_ids'],
        'target_attention_mask': target_encodings['attention_mask']
    }
    print(encodings.keys())
    return encodings

# load train and validation split of squad
#dataset = datasets.load_dataset("Splend1dchan/NMSQA-wav2vecu2")
import pandas as pd
from datasets import Dataset
#./hf/NMSQA-wav2vecu2/
train_dataset = Dataset.from_pandas(pd.read_parquet("./hf/NMSQA-wav2vecu2/train_for_squad_reduce_silIsSpace_dedup.parquet")) #load_dataset()#datasets.load_dataset('Splend1dchan/NMSQA-wav2vecu2', split="train")
valid_dataset = Dataset.from_pandas(pd.read_parquet("./hf/NMSQA-wav2vecu2/dev_for_squad_reduce_silIsSpace_dedup.parquet"))#pd.read_parquet("./hf/NMSQA-wav2vecu2/dev_for_squad_reduce_silIsSpace_dedup.parquet")#datasets.load_dataset('Splend1dchan/NMSQA-wav2vecu2', split="validation")

# map add_eos_to_examples function to the dataset example wise 
train_dataset = train_dataset.map(add_eos_to_examples, load_from_cache_file=False)
# map convert_to_features batch wise
train_dataset = train_dataset.map(convert_to_features, batched=True, load_from_cache_file=False)

valid_dataset = valid_dataset.map(add_eos_to_examples, load_from_cache_file=False)
valid_dataset = valid_dataset.map(convert_to_features, batched=True, load_from_cache_file=False)


# set the tensor type and the columns which the dataset should return
columns = ['input_ids', 'target_ids', 'attention_mask', 'target_attention_mask']
train_dataset.set_format(type="torch", columns=columns)
valid_dataset.set_format(type="torch", columns=columns)
from torch.utils.data import DataLoader
from typing import Dict, List, Optional
def collate_batch(batch: List) -> Dict[str, torch.Tensor]:
    """
    Take a list of samples from a Dataset and collate them into a batch.
    Returns:
        A dictionary of tensors
    """
    #print(batch)
    #s()
    input_ids = torch.stack([example['input_ids'] for example in batch])
    lm_labels = torch.stack([example['target_ids'] for example in batch])
    lm_labels[lm_labels[:, :] == 0] = -100
    attention_mask = torch.stack([example['attention_mask'] for example in batch])
    decoder_attention_mask = torch.stack([example['target_attention_mask'] for example in batch])
    

    return {
        'input_ids': input_ids, 
        'attention_mask': attention_mask,
        'labels': lm_labels, 
        'decoder_attention_mask': decoder_attention_mask
      }
#ld = DataLoader(train_dataset,batch_size = 32,collate_fn = collate_batch)
#for b in ld:
    #print(b)
    #s()
print("training samples",len(train_dataset), "validation samples",len(valid_dataset))

import dataclasses
import logging
import os
import sys
from dataclasses import dataclass, field

import json
import numpy as np
import torch

from transformers import T5ForConditionalGeneration, ByT5Tokenizer, EvalPrediction, T5Tokenizer
from transformers import (
    HfArgumentParser,
    DataCollator,
    Trainer,
    TrainingArguments,
    set_seed,
)

#from tpubar import TPUMonitor

#import torch_xla.core.xla_model as xm
#import torch_xla.distributed.xla_multiprocessing as xmp

logger = logging.getLogger(__name__)

# prepares lm_labels from target_ids, returns examples with keys as expected by the forward method
# this is necessacry because the trainer directly passes this dict as arguments to the model
# so make sure the keys match the parameter names of the forward method
#@dataclass
#class T2TDataCollator(DataCollator):



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

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )

@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
    train_file_path: Optional[str] = field(
        default='train_data.pt',
        metadata={"help": "Path for cached train dataset"},
    )
    valid_file_path: Optional[str] = field(
        default='valid_data.pt',
        metadata={"help": "Path for cached valid dataset"},
    )
    max_len: Optional[int] = field(
        default=512,
        metadata={"help": "Max input length for the source text"},
    )
    target_max_len: Optional[int] = field(
        default=32,
        metadata={"help": "Max input length for the target text"},
    )



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.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))

    # we will load the arguments from a json file, 
    #make sure you save the arguments in at ./args.json
    model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath('args.json'))

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
        )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
    )
    #if xm.is_master_ordinal():
    #logger.warning(
    #    "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
    #    training_args.local_rank,
    #    training_args.device,
    #    training_args.n_gpu,
    #    bool(training_args.local_rank != -1),
    #    training_args.fp16,
    #)
    #logger.info("Training/evaluation parameters %s", training_args)


    # Set seed
    set_seed(training_args.seed)

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    # tokenizer = ByT5Tokenizer.from_pretrained(
    #     model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
    #     cache_dir=model_args.cache_dir,
    # )

    tokenizer = ByT5Tokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )
    config = T5Config.from_pretrained(model_args.model_name_or_path)
    model = T5ForConditionalGeneration(
        config = config
    )
    
    # Get datasets
    #train_dataset  = torch.load(data_args.train_file_path)
    #valid_dataset = torch.load(data_args.valid_file_path)

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        data_collator=collate_batch,
    )

    # Training
    if training_args.do_train:
        trainer.train(
            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
        )
        trainer.save_model()
        # For convenience, we also re-save the tokenizer to the same directory,
        # so that you can share your model easily on huggingface.co/models =)
        if trainer.is_world_master():
            tokenizer.save_pretrained(training_args.output_dir)

    # Evaluation
    results = {}
    if training_args.do_eval and training_args.local_rank in [-1, 0]:
        logger.info("*** Evaluate ***")

        eval_output = trainer.evaluate()

        output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt")
        with open(output_eval_file, "w") as writer:
            logger.info("***** Eval results *****")
            for key in sorted(eval_output.keys()):
                logger.info("  %s = %s", key, str(eval_output[key]))
                writer.write("%s = %s\n" % (key, str(eval_output[key])))
    
        results.update(eval_output)
    
    return results

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()

args_dict = {
  "tpu_num_cores": 1,
  'training_script': 'train_t5_squad.py',
  "model_name_or_path": 'voidful/phoneme_byt5_g2p_v1', #'google/byt5-base',
  "tokenizer_name": 'google/byt5-small', #"google/byt5-base",
  "max_len": 1024 ,
  "target_max_len": 128,
  "output_dir": 'wav2vecu2-byt5small-randinit-squad',
  "overwrite_output_dir": True,
  "per_device_train_batch_size": 4,
  "per_device_eval_batch_size": 1,
  "gradient_accumulation_steps": 4,
  "learning_rate": 3e-4,
  "num_train_epochs": 3,
  "do_train": True,
  "do_eval":False,
  "save_strategy": "epoch",
  "save_total_limit" : 100,
  "push_to_hub" : False,
  "remove_unused_columns" : False
  #"hub_token" : "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"
}
import os
os.environ["WANDB_DISABLED"] = "true"
with open('args.json', 'w') as f:
  json.dump(args_dict, f)
import torch.multiprocessing as mp

if __name__ == '__main__':
    #freeze_support()

    main()
    #mp.spawn(_mp_fn, args=(), nprocs=1, start_method='fork')