training_utils.py

import argparse
from enum import Enum
from transformers import (
    SchedulerType, 
    MODEL_MAPPING
)

from transformers import PreTrainedTokenizerFast, GPT2LMHeadModel
from sklearn.decomposition import PCA
from itertools import chain

from datasets import DatasetDict, load_dataset, concatenate_datasets
import numpy as np
import os

import torch
import math
import pickle

from nltk.translate.bleu_score import sentence_bleu
from nltk.translate.meteor_score import meteor_score
from nltk.tokenize import word_tokenize
from nltk.util import ngrams

MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)

EXACT_MATCH_VALUES = []

def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default=None,
        help="The name of the dataset to use (via the datasets library).",
    )
    parser.add_argument(
        "--dataset_config_name",
        type=str,
        default=None,
        help="The configuration name of the dataset to use (via the datasets library).",
    )
    parser.add_argument(
        "--train_file", type=str, default=None, help="A csv, txt or a json file containing the training data."
    )
    parser.add_argument(
        "--validation_file", type=str, default=None, help="A csv, txt or a json file containing the validation data."
    )
    parser.add_argument(
        "--validation_split_percentage",
        default=5,
        help="The percentage of the train set used as validation set in case there's no validation split",
    )
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
        required=False,
    )
    parser.add_argument(
        "--model_hid_dim",
        type=int,
        help="Hidden dimension of the model.",
        default=None,
        required=False,
    )
    parser.add_argument(
        "--tie_word_embeddings",
        type=int,
        help="Whether to tie the word embeddings with the final layer weights.",
        default=1,
        required=False,
    )
    parser.add_argument(
        "--diff_embs",
        type=int,
        help="Whether to use different dimensions for embedding matrix and model hidden dimension.",
        default=0,
        required=False,
    )
    
    parser.add_argument(
        "--frac_seq_inc_freq",
        type=float,
        help="Fraction of sequences to increase the frequency of while constructing the data.",
        default=0,
        required=False,
    )
    
    parser.add_argument(
        "--inc_freq_mult",
        type=int,
        help="Multiplier to use for increasing the frequency of chosen frac_seq_inc_freq \% sequences. For instance, a value of 2 means that the chosen sequences appear twice as often in the training data.",
        default=1,
        required=False,
    )
    
    parser.add_argument(
        "--config_name",
        type=str,
        default=None,
        help="Pretrained config name or path if not the same as model_name",
    )
    parser.add_argument(
        "--tokenizer_name",
        type=str,
        default=None,
        help="Pretrained tokenizer name or path if not the same as model_name",
    )
    parser.add_argument(
        "--use_slow_tokenizer",
        action="store_true",
        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
    )
    parser.add_argument(
        "--sequence_shuffle",
        action="store_true",
        help="If passed, will shuffle the words in each sequence before training.",
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the evaluation dataloader.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-5,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )

    parser.add_argument(
        "--grouping_strategy",
        type=str,
        default="fit_max",
        help=(
            "What grouping strategy to use? It could be one of the following:"
            "(1) 'fit_max' - fit as many examples as possible in each training sequence, pad the remaining tokens until block size."
            "(2) a number - denoting how many rollouts must be fitted into each training sequence. Note that context length is assumed to be sufficient to fit the asked number of rollouts."
        )
    )

    parser.add_argument(
        "--random_data_frac",
        type=float,
        default=0,
        help="Fraction of random data to use for training. \
            For a value k, original data and random data are mixed in proportions (1-k) and k respectively. \
            The random data used is the sequences of random digits from 0 to 9."
    )
    parser.add_argument(
        "--padding_side",
        type=str,
        default="right",
        help="The side on which to pad the input sequences.",
        choices=["left", "right"],
    )

    parser.add_argument(
        "--lr_scheduler_type",
        type=str,
        default="linear",
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup", "lambdaLR"],
    )

    parser.add_argument(
        "--dec_lr_epoch",
        type=int,
        default=None,
        help="The epoch at which to decrease the learning rate to one-tenth."
    )

    parser.add_argument(
        "--wandb_project_name",
        type=str,
        default=None,
        help="Name of the wandb project to log"
    )

    parser.add_argument(
        "--wandb_entity",
        type=str,
        default=None,
        help="Entity name for wandb logging."
    )

    parser.add_argument(
        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument(
        "--model_type",
        type=str,
        default=None,
        help="Model type to use if training from scratch.",
        choices=MODEL_TYPES,
    )
    parser.add_argument(
        "--dropout",
        type=float,
        default=0.1,
        help="Dropout to use in the model. The specified fraction of neurons are randomly dropped."
    )

    parser.add_argument(
        "--model_n_layer",
        type=int,
        default=None,
        help="Number of layers in the model."
    )

    parser.add_argument(
        "--model_n_head",
        type=int,
        default=None,
        help="Number of attention heads in the model."
    )

    parser.add_argument(
        "--model_n_embd",
        type=int,
        default=None,
        help="Size of model's embedding. It is used as the model hidden dim too if diff_embs is 0."
    )

    parser.add_argument(
        "--num_params",
        type=float,
        default=None,
        help="Number of parameters (in M) in the model."
    )

    parser.add_argument(
        "--model_index",
        type=int,
        default=None,
        help="Index of the model config to train on in the list of models with the same number of parameters."
    )

    parser.add_argument(
        "--data_fraction",
        type=float,
        default=None,
        help="Fraction of data to use for training."
    )

    parser.add_argument(
        "--save_name",
        type=str,
        default=None,
        help="Model is saved at the path {output_dir}/{save_name}. {save_name} is also used as the run name for tracking in wandb."
    )

    parser.add_argument(
        "--block_size",
        type=int,
        default=None,
        help=(
            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
            " this size for training. Default to the model max input length for single sentence inputs (take into"
            " account special tokens)."
        ),
    )

    parser.add_argument(
        "--freeze_emb",
        type=str,
        default=None,
        choices=[None, "random_init", "gpt2_init"],
        help="Whether to freeze the token and position embeddings. Options:\
        (1) None - don't freeze any embeddings,\
        (2) 'random_init' - freeze the embeddings and initialize them randomly,\
        (3) 'gpt2_init' - freeze the embeddings and initialize them with GPT2 embeddings."
    )
    
    parser.add_argument(
        "--print_gen_steps",
        type=int,
        default=1000,
        help="Print generations after every print_gen_steps steps."
    )

    parser.add_argument(
        "--preprocessing_num_workers",
        type=int,
        default=None,
        help="The number of processes to use for the preprocessing.",
    )
    parser.add_argument(
        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
    )
    parser.add_argument(
        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
    )
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument(
        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
    )
    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
    parser.add_argument(
        "--trust_remote_code",
        action="store_true",
        help=(
            "Whether to trust the execution of code from datasets/models defined on the Hub."
            " This option should only be set to `True` for repositories you trust and in which you have read the"
            " code, as it will execute code present on the Hub on your local machine."
        ),
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=str,
        default=None,
        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
    )

    parser.add_argument(
        "--save_every_epoch",
        type=str,
        default="total_5",
        help="Epoch saving strategy. Only considered if 'checkpointing_steps' = 'epoch'. Options: \
        [\
            'total_k' - save a total of k ckpts uniformly distributed across training if 'num_train_epochs' > k else save every epoch, \
            'every_k' (int) - save ckpt every k epochs if 'num_train_epochs' > k else save every epoch, \
            'save-only-last' - save only the final ckpt (i.e. at the end of training), \
            'no-ckpt' - don't save any ckpt, \
        ]. Final ckpt is always saved for every option except 'no-ckpt'."
    )

    parser.add_argument(
        "--eval_every_steps",
        type=int,
        default=1000,
        help="Perform evaluation after every eval_every_steps train steps and log the evaluation results.",
    )

    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help="If the training should continue from a checkpoint folder.",
    )
    parser.add_argument(
        "--with_tracking",
        action="store_true",
        help="Whether to enable experiment trackers for logging.",
    )
    parser.add_argument(
        "--report_to",
        type=str,
        default="all",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
            ' `"wandb"`, `"comet_ml"` and `"clearml"`. Use `"all"` (default) to report to all integrations. '
            "Only applicable when `--with_tracking` is passed."
        ),
    )
    parser.add_argument(
        "--low_cpu_mem_usage",
        action="store_true",
        help=(
            "It is an option to create the model as an empty shell, then only materialize its parameters when the pretrained weights are loaded. "
            "If passed, LLM loading time and RAM consumption will be benefited."
        ),
    )
    args = parser.parse_args()

    # Sanity checks
    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
        raise ValueError("Need either a dataset name or a training/validation file.")
    else:
        if args.train_file is not None:
            extension = args.train_file.split(".")[-1]
            if extension not in ["csv", "json", "txt"]:
                raise ValueError("`train_file` should be a csv, json or txt file.")
        if args.validation_file is not None:
            extension = args.validation_file.split(".")[-1]
            if extension not in ["csv", "json", "txt"]:
                raise ValueError("`validation_file` should be a csv, json or txt file.")

    if args.push_to_hub:
        if args.output_dir is None:
            raise ValueError("Need an `output_dir` to create a repo when `--push_to_hub` is passed.")

    # check that print_gen_steps is a multiple of eval_every_steps
    if args.print_gen_steps % args.eval_every_steps != 0:
        raise ValueError("print_gen_steps must be a multiple of eval_every_steps.")
    
    return args

def group_texts(examples, tokenizer: PreTrainedTokenizerFast, block_size, grouping_strategy=1, padding_side="left", loss_type="ATL"):
    len_ex_inp_ids = len(examples['input_ids'])

    # grouping
    tokens_remaining = block_size
    grouped_token_ids = []
    current_group = []
    if grouping_strategy == 'fit_max':
        for idx in range(len_ex_inp_ids):
            input_id_list = examples['input_ids'][idx]
            if len(input_id_list) > block_size:
                continue
            if len(input_id_list) <= tokens_remaining:
                current_group.append(input_id_list)
                tokens_remaining -= len(input_id_list)
                if idx == len_ex_inp_ids-1:
                    current_group_flat = list(chain(*current_group))
                    if padding_side == "left":
                        current_group_flat = [tokenizer.pad_token_id]*(block_size-len(current_group_flat)) + current_group_flat
                    else:
                        current_group_flat = current_group_flat + [tokenizer.pad_token_id]*(block_size-len(current_group_flat))
                    grouped_token_ids.append(current_group_flat)
                    current_group = []
                    tokens_remaining = block_size
            else:
                current_group_flat = list(chain(*current_group))
                if padding_side == "left":
                    current_group_flat = [tokenizer.pad_token_id]*(block_size-len(current_group_flat)) + current_group_flat
                else:
                    current_group_flat = current_group_flat + [tokenizer.pad_token_id]*(block_size-len(current_group_flat))
                grouped_token_ids.append(current_group_flat)
                current_group = []
                tokens_remaining = block_size

                current_group.append(input_id_list)
                tokens_remaining -= len(input_id_list)
                if idx == len_ex_inp_ids-1:
                    current_group_flat = list(chain(*current_group))
                    if padding_side == "left":
                        current_group_flat = [tokenizer.pad_token_id]*(block_size-len(current_group_flat)) + current_group_flat
                    else:
                        current_group_flat = current_group_flat + [tokenizer.pad_token_id]*(block_size-len(current_group_flat))
                    grouped_token_ids.append(current_group_flat)
                    current_group = []
                    tokens_remaining = block_size
    elif isinstance(grouping_strategy, int):
        for group_start_idx in range(0, len_ex_inp_ids, grouping_strategy):
            input_id_lists = examples['input_ids'][group_start_idx: (group_start_idx + grouping_strategy)]
            input_id_lists_flatten = list(chain(*input_id_lists))
            if padding_side == "left":
                current_group_flat = [tokenizer.pad_token_id]*(block_size-len(input_id_lists_flatten)) + input_id_lists_flatten
            else:
                current_group_flat = input_id_lists_flatten + [tokenizer.pad_token_id]*(block_size-len(input_id_lists_flatten))
            grouped_token_ids.append(current_group_flat)

    attention_masks = []
    labels = []

    for group in grouped_token_ids:
        assert len(group) == block_size
        attention_masks.append([0 if tkn_id == tokenizer.pad_token_id else 1 for tkn_id in group])
        
        if loss_type == "ATL":
            labels.append([-100 if tkn_id == tokenizer.pad_token_id else tkn_id for tkn_id in group])
        else:
            raise NotImplementedError

    result = {
        "input_ids": grouped_token_ids,
        "attention_mask": attention_masks,
        "labels": labels,
    }
    
    return result

def get_epoch_save_steps(save_every_epoch, num_train_epochs):
    epoch_save_steps = None
    if save_every_epoch.startswith("total"):
        total_ckpt_num = int(save_every_epoch.split("_")[1])
        epoch_save_steps = num_train_epochs // total_ckpt_num if num_train_epochs > total_ckpt_num else 1
    elif save_every_epoch.startswith("every"):
        save_num = int(save_every_epoch.split("_")[1])
        epoch_save_steps = save_num if num_train_epochs > save_num else int(1e9)
    elif save_every_epoch == 'save-only-last':
        epoch_save_steps = int(1e9) # very large value so that when we compute modulus with this divisor, the dividend itself is the remainder; 
                                # we'll never get a zero divisor and hence will never be saving ckpt
    elif save_every_epoch == 'no-ckpt':
        epoch_save_steps = None
    else:
        raise NotImplementedError
    return epoch_save_steps

def num_params_to_model_params(args):
    """
        Convert number of parameters to model parameters
    """
    num_params, model_index = args.num_params, args.model_index

    # num_params: (L, E, H) -> (n_layer, n_embd, n_head)
    num_params_to_model_params_dict = {
        0.5: (1, 8, 2),
        1: (1, 16, 4),
        2: (1, 32, 4),
        4: (1, 64, 8),
        8: (1, 128, 16),
        16: (1, 256, 16),
        32: (1, 512, 16),
        40: (1, 768, 16),
    }

    num_params_to_model_params_dict[0.5] = [
        (2, 144, 1),
        (1, 10, 1),
        (16, 10, 1)
    ]

    num_params_to_model_params_dict[1] = [
        (2, 204, 1),
        (1, 20, 1),
        (16, 20, 1)
    ]

    num_params_to_model_params_dict[2] = [
        (3, 235, 1),
        (2, 40, 1),
    ]

    num_params_to_model_params_dict[4] = [
        (3, 333, 1),
        (4, 76, 1),
    ]

    num_params_to_model_params_dict[8] = [
        (4, 408, 1),
        (8, 128, 1),
    ]

    num_params_to_model_params_dict[12] = [
        (4, 500, 1),
    ]

    num_params_to_model_params_dict[16] = [
        (5, 516, 1),
        (1, 292, 1),
        (1, 300, 1),
        (8, 216, 1),
        (8, 224, 1),
        (16, 184, 1),
        (16, 188, 1),
        (24, 160, 1),
        (24, 164, 1),
        (32, 148, 1),
        (40, 136, 1),
        (48, 128, 1)
    ]

    num_params_to_model_params_dict[20] = [
        (5, 577, 1),
    ]

    num_params_to_model_params_dict[24] = [
        (5, 632, 1),
    ]

    num_params_to_model_params_dict[28] = [
        (5, 683, 1),
    ]

    num_params_to_model_params_dict[32] = [
        (6, 666, 1),
        (1, 552, 1),
        (1, 560, 1),
        (8, 368, 1),
        (8, 372, 1),
        (16, 296, 1),
        (24, 256, 1),
        (32, 228, 1),
        (40, 208, 1),
        (48, 196, 1)
    ]

    num_params_to_model_params_dict[124] = [
        (12, 768, 12),
    ]

    n_layer, n_embd, n_head = num_params_to_model_params_dict[num_params][model_index]

    if args.model_n_layer is not None and args.model_n_embd is not None and args.model_n_head is not None:
        n_layer, n_embd, n_head = args.model_n_layer, args.model_n_embd, args.model_n_head

    if args.diff_embs == 1:
        assert args.model_hid_dim is not None, "Model hidden dimension is required when args.diff_embs is 1 (using different dimensions for embedding matrix and model hidden dimension.)"
        model_hid_dim = args.model_hid_dim
    else:
        model_hid_dim = n_embd
    
    return n_layer, n_embd, model_hid_dim, n_head

# Function to get frac% of the dataset and store indices
def get_fraction_one_dataset(dataset, frac, random_dataset=None, random_data_frac=0, split=None, args=None):
    # Calculate the number of samples to keep
    if random_dataset is None:
        if split == "train":
            num_samples = int(frac * len(dataset))
            
            num_seq_repeat = int(args.frac_seq_inc_freq * num_samples)
            num_non_repeating_seq = num_samples - args.inc_freq_mult * num_seq_repeat
            
            # Generate random indices
            total_seq_select = num_non_repeating_seq + num_seq_repeat
            indices = np.random.choice(len(dataset), total_seq_select, replace=False)
            
            # Select the samples
            new_dataset1 = dataset.select(indices[:num_seq_repeat]) # first num_seq_repeat indices are for the sequences to be repeated
            new_dataset2 = dataset.select(indices[num_seq_repeat:]) # next num_non_repeating_seq indices are for the sequences not to be repeated
            
            new_dataset1_rep = concatenate_datasets([new_dataset1] * args.inc_freq_mult)
            
            new_dataset = concatenate_datasets([new_dataset1_rep, new_dataset2])
        else:
            num_samples = int(frac * len(dataset))
            indices = np.random.choice(len(dataset), num_samples, replace=False)
            new_dataset = dataset.select(indices)
    else:
        num_samples = int(frac * len(dataset))
        num_random_samples = int(random_data_frac * num_samples)
        num_original_samples = num_samples - num_random_samples

        # Generate random indices
        indices = np.random.choice(len(dataset), num_original_samples, replace=False)
        random_indices = np.random.choice(len(random_dataset), num_random_samples, replace=False)

        # Select the samples
        new_dataset = dataset.select(indices)
        random_new_dataset = random_dataset.select(random_indices)
        new_dataset = concatenate_datasets([new_dataset, random_new_dataset])
    
    return new_dataset, indices

def make_fractional_dataset_dict(original_dataset_dict, data_fraction, output_dir, random_dataset=None, random_data_frac=0, args=None):
    """
        Make a fractional dataset
    """
    # Create new DatasetDict and store indices
    new_dataset_dict = DatasetDict()
    indices_dict = {}
    
    assert (random_dataset is None) or (args.frac_seq_inc_freq == 0), "The features for adding random data and to increase frequency of sequences cannot be used together."

    for split in original_dataset_dict:
        random_dataset_split = split
        if split == "test":
            random_dataset_split = "validation" # random_dataset does not have a test split
        if random_dataset is not None:
            new_dataset, indices = get_fraction_one_dataset(original_dataset_dict[split], data_fraction, random_dataset=random_dataset[random_dataset_split], random_data_frac=random_data_frac)
        else:
            new_dataset, indices = get_fraction_one_dataset(original_dataset_dict[split], data_fraction, split=split, args=args)
        new_dataset_dict[split] = new_dataset
        indices_dict[split] = indices

    # Save the indices_dict for later reconstruction
    for split, indices in indices_dict.items():
        os.makedirs(os.path.join(output_dir, "selected_indices"), exist_ok=True)
        path = os.path.join(output_dir, "selected_indices", f'selected_indices_{split}.npy')
        np.save(path, indices)
    
    return new_dataset_dict

def load_npy_and_print_indices_shape(output_dir, splits):
    for split in splits:
        path = os.path.join(output_dir, "selected_indices", f'selected_indices_{split}.npy')
        indices = np.load(path)
        print(f'Shape of indices for {split}: {indices.shape}')


def do_eval_and_log(model, eval_dataloader, args, accelerator, epoch, completed_steps, logger, total_loss, avg_interval_loss, train_dataloader, train_dataset, tokenizer, eval_full_train=False, print_gen=False, memories=None, unmem_seq_metrics=None, prompt_seq2idx=None, gen_seq2idx=None):
    model.eval()

    if args.data_fraction >= 0.05:
        losses = []
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                outputs = model(**batch)

            loss = outputs.loss
            losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))

        losses = torch.cat(losses)
        try:
            eval_loss = torch.mean(losses)
            perplexity = math.exp(eval_loss)
        except OverflowError:
            perplexity = float("inf")

        logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
    else:
        perplexity=-1
        eval_loss=-1
        # logger.info(f"epoch {epoch}: train ppl: {math.exp(total_loss.item() / len(train_dataloader))} train_loss: {total_loss.item() / len(train_dataloader)}")
        logger.info(f"epoch {epoch}: train ppl: {math.exp(avg_interval_loss.item())} train_loss: {avg_interval_loss.item()}")

    bleu, meteor, exact_match, jaccard_similarity, set_mem_seq, unmem_metrics, bleu_scores_10tok, overall_meteor_score_10tok, overall_exact_match_10tok = do_generation_eval(tokenizer, train_dataset, model, eval_full_train=eval_full_train, print_gen=print_gen, prompt_seq2idx=prompt_seq2idx, gen_seq2idx=gen_seq2idx)

    EXACT_MATCH_VALUES.append((completed_steps, exact_match))
    
    memories[completed_steps] = set_mem_seq
    unmem_seq_metrics[completed_steps] = unmem_metrics
    
    if args.with_tracking:
        accelerator.log(
            {
                "val_ppl": perplexity,
                "val_loss": eval_loss,
                "train_loss": avg_interval_loss.item(),
                "train_ppl": math.exp(avg_interval_loss.item()),
                "bleu_indiv-1": bleu['indiv 1-gram'],
                "bleu_indiv-2": bleu['indiv 2-gram'],
                "bleu_indiv-3": bleu['indiv 3-gram'],
                "bleu_indiv-4": bleu['indiv 4-gram'],
                "bleu_indiv-5": bleu['indiv 5-gram'],
                "bleu_indiv-10": bleu['indiv 10-gram'],
                "bleu_indiv-15": bleu['indiv 15-gram'],
                "bleu_indiv-20": bleu['indiv 20-gram'],
                "bleu_cum-1": bleu['cumulative 1-gram'],
                "bleu_cum-2": bleu['cumulative 2-gram'],
                "bleu_cum-3": bleu['cumulative 3-gram'],
                "bleu_cum-4": bleu['cumulative 4-gram'],
                "exact_match": exact_match,
                "jaccard_similarity": jaccard_similarity,
                "meteor": meteor,
                "bleu_10tok_indiv-1": bleu_scores_10tok['indiv 1-gram'],
                "bleu_10tok_indiv-2": bleu_scores_10tok['indiv 2-gram'],
                "bleu_10tok_indiv-3": bleu_scores_10tok['indiv 3-gram'],
                "bleu_10tok_indiv-4": bleu_scores_10tok['indiv 4-gram'],
                "bleu_10tok_cum-1": bleu_scores_10tok['cumulative 1-gram'],
                "bleu_10tok_cum-2": bleu_scores_10tok['cumulative 2-gram'],
                "bleu_10tok_cum-3": bleu_scores_10tok['cumulative 3-gram'],
                "bleu_10tok_cum-4": bleu_scores_10tok['cumulative 4-gram'],
                "meteor_10tok": overall_meteor_score_10tok,
                "exact_match_10tok": overall_exact_match_10tok,
                "epoch": epoch,
                "step": completed_steps,
            },
            step=completed_steps,
        )
    
    model.train()
    return eval_loss, perplexity

def get_ngrams(tokenized_text, n):
    return set(ngrams(tokenized_text, n))

def jaccard_similarity(set1, set2):
    intersection = set1.intersection(set2)
    union = set1.union(set2)
    return len(intersection) / len(union)

def calculate_all_metrics(text, output, bleu_scores, jacc_ngram=13, no_jacc=False, tok_10=False):
    # Calculate BLEU scores
    tokenized_text = word_tokenize(text)
    tokenized_output = word_tokenize(output)
    
    meteor_score_ = meteor_score([tokenized_text], tokenized_output)
    
    if not no_jacc:
        # Generate 13-grams
        ngrams1 = get_ngrams(tokenized_text, jacc_ngram)
        ngrams2 = get_ngrams(tokenized_output, jacc_ngram)
        if len(ngrams1) == 0 and len(ngrams2) == 0:
            jacc_similarity = None
        else:
            # Compute Jaccard similarity
            jacc_similarity = jaccard_similarity(ngrams1, ngrams2)
    
    bleu_indiv_score_1 = sentence_bleu([tokenized_text], tokenized_output, weights=(1, 0, 0, 0))
    bleu_indiv_score_2 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 1, 0, 0))
    bleu_indiv_score_3 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 0, 1, 0))
    bleu_indiv_score_4 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 0, 0, 1))
    if not tok_10:
        bleu_indiv_score_5 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 0, 0, 0, 1))
        bleu_indiv_score_10 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 0, 0, 0, 0, 0, 0, 0, 0, 1))
        bleu_indiv_score_15 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1))
        bleu_indiv_score_20 = sentence_bleu([tokenized_text], tokenized_output, weights=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1))

    bleu_scores["indiv 1-gram"].append(bleu_indiv_score_1)
    bleu_scores["indiv 2-gram"].append(bleu_indiv_score_2)
    bleu_scores["indiv 3-gram"].append(bleu_indiv_score_3)
    bleu_scores["indiv 4-gram"].append(bleu_indiv_score_4)
    if not tok_10:
        bleu_scores["indiv 5-gram"].append(bleu_indiv_score_5)
        bleu_scores["indiv 10-gram"].append(bleu_indiv_score_10)
        bleu_scores["indiv 15-gram"].append(bleu_indiv_score_15)
        bleu_scores["indiv 20-gram"].append(bleu_indiv_score_20)
    
    bleu_score_cum_1 = sentence_bleu([tokenized_text], tokenized_output, weights=(1, 0, 0, 0))
    bleu_score_cum_2 = sentence_bleu([tokenized_text], tokenized_output, weights=(0.5, 0.5, 0, 0))
    bleu_score_cum_3 = sentence_bleu([tokenized_text], tokenized_output, weights=(0.33, 0.33, 0.33, 0))
    bleu_score_cum_4 = sentence_bleu([tokenized_text], tokenized_output, weights=(0.25, 0.25, 0.25, 0.25))
    
    bleu_scores["cumulative 1-gram"].append(bleu_score_cum_1)
    bleu_scores["cumulative 2-gram"].append(bleu_score_cum_2)
    bleu_scores["cumulative 3-gram"].append(bleu_score_cum_3)
    bleu_scores["cumulative 4-gram"].append(bleu_score_cum_4)
    
    if no_jacc:
        return meteor_score_, None
    return meteor_score_, jacc_similarity

NUM_GENERATIONS_PRINT = 5
EVAL_NUM_EX = 90 # 1000
PROMPT_LEN = 30
GENERATE_LEN = 70

def do_generation_eval(tokenizer, train_dataset, model, eval_full_train=False, print_gen=False, prompt_seq2idx=None, gen_seq2idx=None):
    PRINT_MODEL_GENERATIONS = print_gen
    
    print("Generating and Evaluating...")
    # calculate the BLEU and METEOR scores
    bleu_scores = {"indiv 1-gram": [], "indiv 2-gram": [], "indiv 3-gram": [], "indiv 4-gram": [], 
                     "indiv 5-gram": [], "indiv 10-gram": [], "indiv 15-gram": [], "indiv 20-gram": [],
                   "cumulative 1-gram": [], "cumulative 2-gram": [], "cumulative 3-gram": [], "cumulative 4-gram": []}
    meteor_scores = []
    jacc_similarities = []
    exact_match = 0
    set_mem_seq = set()
    unmem_metrics = {}
    
    # 10-tok metrics
    bleu_scores_10tok = {"indiv 1-gram": [], "indiv 2-gram": [], "indiv 3-gram": [], "indiv 4-gram": [], 
                   "cumulative 1-gram": [], "cumulative 2-gram": [], "cumulative 3-gram": [], "cumulative 4-gram": []}
    meteor_scores_10tok = []
    # jacc_similarities_10tok = []
    exact_match_10tok = 0

    if eval_full_train:
        print("Evaluating on full train dataset...")
        train = train_dataset
    else:
        random_indices = np.random.choice(len(train_dataset), EVAL_NUM_EX, replace=False)
        train = train_dataset.select(random_indices)

    # Prepare the input tokens for all examples at once
    total_len = PROMPT_LEN + GENERATE_LEN

    input_tokens = {
            'input_ids': torch.cat([torch.tensor(batch['input_ids']).view(1, -1) for batch in train], dim=0).to('cuda'),
            'attention_mask': torch.cat([torch.tensor(batch['attention_mask']).view(1, -1) for batch in train], dim=0).to('cuda'),
            'labels': torch.cat([torch.tensor(batch['labels']).view(1, -1) for batch in train], dim=0).to('cuda')
        }

    input_tokens_prompt_ids = input_tokens['input_ids'][:, :PROMPT_LEN]
    input_tokens_prompt_att_mask = input_tokens['attention_mask'][:, :PROMPT_LEN]

    input_tokens_generate_ids_ground_truth = input_tokens['input_ids'][:, PROMPT_LEN: total_len]
    input_tokens_generate_ids_ground_truth_10_tok = input_tokens['input_ids'][:, PROMPT_LEN: (PROMPT_LEN + 10)]

    model.eval()
    # Generate outputs for all examples at once
    output_ids = model.generate(
        input_tokens_prompt_ids, 
        attention_mask=input_tokens_prompt_att_mask,
        max_length=total_len, 
        num_return_sequences=1, 
        pad_token_id=tokenizer.pad_token_id,
        do_sample=False,
    )

    generated_token_ids = output_ids[:, input_tokens_prompt_ids.shape[1]:]

    prompt = tokenizer.batch_decode(input_tokens_prompt_ids, skip_special_tokens=True)
    model_generated = tokenizer.batch_decode(generated_token_ids, skip_special_tokens=True)
    model_generated_10tok = tokenizer.batch_decode(generated_token_ids[:, :10], skip_special_tokens=True)
    ground_truth = tokenizer.batch_decode(input_tokens_generate_ids_ground_truth, skip_special_tokens=True)
    ground_truth_10tok = tokenizer.batch_decode(input_tokens_generate_ids_ground_truth_10_tok, skip_special_tokens=True)

    # breakpoint()
    if PRINT_MODEL_GENERATIONS:
        # select a few random examples to print by indexing into the model_generated and ground_truth lists with NUM_GENERATIONS_PRINT random indices
        random_indices = np.random.choice(EVAL_NUM_EX, NUM_GENERATIONS_PRINT, replace=False)
        subset_model_generated = [model_generated[i] for i in random_indices]
        subset_ground_truth = [ground_truth[i] for i in random_indices]

        for i in range(NUM_GENERATIONS_PRINT):
            print(f"Example {i}")
            print(f"Prompt: {prompt[random_indices[i]]}")
            print(f"Model generated: {subset_model_generated[i]}")
            print(f"Ground truth: {subset_ground_truth[i]}")
            print()

    # Calculate the BLEU and METEOR scores between model generated and ground truth texts
    
    for i, text, text_10tok in zip(range(len(ground_truth)), ground_truth, ground_truth_10tok):
        output = model_generated[i]
        meteor_score_, jacc_similarity = calculate_all_metrics(text, output, bleu_scores)
        meteor_scores.append(meteor_score_)
        if jacc_similarity is not None:
            jacc_similarities.append(jacc_similarity)
        if text == output:
            exact_match += 1
            set_mem_seq.add(gen_seq2idx[text])
        else:
            bleu_scores_for_this_seq = {
                "indiv 1-gram": bleu_scores["indiv 1-gram"][-1],
                "indiv 2-gram": bleu_scores["indiv 2-gram"][-1],
                "indiv 3-gram": bleu_scores["indiv 3-gram"][-1],
                "indiv 4-gram": bleu_scores["indiv 4-gram"][-1],
                "indiv 5-gram": bleu_scores["indiv 5-gram"][-1],
                "indiv 10-gram": bleu_scores["indiv 10-gram"][-1],
                "indiv 15-gram": bleu_scores["indiv 15-gram"][-1],
                "indiv 20-gram": bleu_scores["indiv 20-gram"][-1],
                "cumulative 1-gram": bleu_scores["cumulative 1-gram"][-1],
                "cumulative 2-gram": bleu_scores["cumulative 2-gram"][-1],
                "cumulative 3-gram": bleu_scores["cumulative 3-gram"][-1],
                "cumulative 4-gram": bleu_scores["cumulative 4-gram"][-1]
            }
            unmem_metrics[i] = {**bleu_scores_for_this_seq, "meteor": meteor_score_, "jaccard_similarity": jacc_similarity}
        
        output_10tok = model_generated_10tok[i]
        meteor_score_10tok, _ = calculate_all_metrics(text_10tok, output_10tok, bleu_scores_10tok, no_jacc=True, tok_10=True)
        meteor_scores_10tok.append(meteor_score_10tok)
        
        if text_10tok == output_10tok:
            exact_match_10tok += 1

    # average all the bleu scores in dict belu_scores
    for key in bleu_scores:
        bleu_scores[key] = sum(bleu_scores[key]) / len(bleu_scores[key])
        
    for key in bleu_scores_10tok:
        bleu_scores_10tok[key] = sum(bleu_scores_10tok[key]) / len(bleu_scores_10tok[key])
    
    overall_meteor_score = sum(meteor_scores) / len(meteor_scores)
    overall_jacc_similarity = sum(jacc_similarities) / len(jacc_similarities)
    overall_exact_match = exact_match / len(ground_truth)
    
    overall_meteor_score_10tok = sum(meteor_scores_10tok) / len(meteor_scores_10tok)
    overall_exact_match_10tok = exact_match_10tok / len(ground_truth_10tok)
    print('Individual BLEU scores:')
    print('Individual 1-gram:', bleu_scores['indiv 1-gram'])
    print('Individual 2-gram:', bleu_scores['indiv 2-gram'])
    print('Individual 3-gram:', bleu_scores['indiv 3-gram'])
    print('Individual 4-gram:', bleu_scores['indiv 4-gram'])
    print('Individual 5-gram:', bleu_scores['indiv 5-gram'])
    print('Individual 10-gram:', bleu_scores['indiv 10-gram'])
    print('Individual 15-gram:', bleu_scores['indiv 15-gram'])
    print('Individual 20-gram:', bleu_scores['indiv 20-gram'])
    print('Cumulative BLEU scores:')
    print('BLEU-1:', bleu_scores['cumulative 1-gram'])
    print('BLEU-2:', bleu_scores['cumulative 2-gram'])
    print('BLEU-3:', bleu_scores['cumulative 3-gram'])
    print('BLEU-4:', bleu_scores['cumulative 4-gram'])
    print('Overall METEOR score:', overall_meteor_score)
    print('Overall Exact match:', overall_exact_match)
    print('Overall 13-gram Jaccard similarity:', overall_jacc_similarity)
    
    print()
    print("10-tok metrics")
    print('Individual BLEU scores:')
    print('Individual 1-gram:', bleu_scores_10tok['indiv 1-gram'])
    print('Individual 2-gram:', bleu_scores_10tok['indiv 2-gram'])
    print('Individual 3-gram:', bleu_scores_10tok['indiv 3-gram'])
    print('Individual 4-gram:', bleu_scores_10tok['indiv 4-gram'])
    print('Cumulative BLEU scores:')
    print('BLEU-1:', bleu_scores_10tok['cumulative 1-gram'])
    print('BLEU-2:', bleu_scores_10tok['cumulative 2-gram'])
    print('BLEU-3:', bleu_scores_10tok['cumulative 3-gram'])
    print('BLEU-4:', bleu_scores_10tok['cumulative 4-gram'])
    print('Overall METEOR score:', overall_meteor_score_10tok)
    print('Overall Exact match:', overall_exact_match_10tok)

    model.train()
    return bleu_scores, overall_meteor_score, overall_exact_match, overall_jacc_similarity, set_mem_seq, unmem_metrics, bleu_scores_10tok, overall_meteor_score_10tok, overall_exact_match_10tok
    
def get_size(L, E, n_pos, vocab_size):
    emb_params = E * vocab_size
    position_emb_params = E * n_pos
    total_emb_params = emb_params + position_emb_params

    block_params_attention = 4 * E**2 + 4 * E # 4d^2 for Q, K, V, O and 4d for bias
    block_params_MLP = 8 * E**2 + 5 * E # d to 4d and then 4d to d => 8d^2; 5d for bias
    block_params_ln = 2 * (2 * E) # 2 layer norms per layer with 2E parameters each for scale and bias
    block_params_total = block_params_attention + block_params_MLP + block_params_ln
    all_blocks_params = L * block_params_total

    final_ln_params = 2 * E

    total_params_downstream = all_blocks_params + final_ln_params
    total_params = total_emb_params + total_params_downstream

    total_params_2_dec = round(total_params/1000**2, 2)
    total_params_downstream_2_dec = round(total_params_downstream/1000**2, 2)
    total_emb_params_2_dec = round(total_emb_params/1000**2, 2)

    return total_params_2_dec, total_params_downstream_2_dec, total_emb_params_2_dec

def get_gpt2_based_on_emb_dim(emb_dim):
    if emb_dim <= 124:
        return "gpt2"
    elif emb_dim <= 355:
        return "gpt2-medium"
    elif emb_dim <= 774:
        return "gpt2-large"
    elif emb_dim <= 1542:
        return "gpt2-xl"
    else:
        raise NotImplementedError

def init_with_gpt2_emb(model, args):
    emb_dim = model.config.n_embd
    gpt2_model_name = get_gpt2_based_on_emb_dim(emb_dim)
    print(f"Initializing with GPT2 embeddings from {gpt2_model_name}...")
    gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
    
    # get model's token and position embeddings
    tok_emb = gpt2_model.transformer.wte.weight
    pos_emb = gpt2_model.transformer.wpe.weight

    # Reduce the dimensionality of the token embeddings using PCA
    pca_tok = PCA(n_components=emb_dim)
    tok_emb_pca = pca_tok.fit_transform(tok_emb.detach().numpy())

    # Reduce the dimensionality of the position embeddings using PCA
    pca_pos = PCA(n_components=emb_dim)
    pos_emb_pca = pca_pos.fit_transform(pos_emb.detach().numpy())

    # Create PyTorch tensors from the PCA-transformed embeddings
    tok_emb_pca_tensor = torch.tensor(tok_emb_pca, dtype=tok_emb.dtype)
    pos_emb_pca_tensor = torch.tensor(pos_emb_pca, dtype=pos_emb.dtype)
    
    # Copy the PCA-transformed embeddings to the model
    model.transformer.wte.weight.data = tok_emb_pca_tensor
    model.transformer.wpe.weight.data = pos_emb_pca_tensor[:args.block_size, :]

def dump_seq2idx_mapping(train_dataset, tokenizer, output_dir):
    train = train_dataset

    # Prepare the input tokens for all examples at once
    total_len = PROMPT_LEN + GENERATE_LEN

    input_tokens = {
        'input_ids': torch.cat([torch.tensor(batch['input_ids']).view(1, -1) for batch in train], dim=0).to('cuda')
    }
    
    input_tokens_prompt_ids = input_tokens['input_ids'][:, :PROMPT_LEN]
    prompt = tokenizer.batch_decode(input_tokens_prompt_ids, skip_special_tokens=True)

    input_tokens_generate_ids_ground_truth = input_tokens['input_ids'][:, PROMPT_LEN: total_len]
    ground_truth = tokenizer.batch_decode(input_tokens_generate_ids_ground_truth, skip_special_tokens=True)

    prompt_seq2idx = {}
    gen_seq2idx = {}
    
    for i, pseq, gseq in zip(range(len(prompt)), prompt, ground_truth):
        prompt_seq2idx[pseq] = i
        gen_seq2idx[gseq] = i
        
    os.makedirs(os.path.join(output_dir, "seq2idx_mapping"), exist_ok=True)
    path = os.path.join(output_dir, "seq2idx_mapping", f'prompt_seq2idx.pkl')
    
    # pickle dump the seq2idx mappings
    with open(path, 'wb') as f:
        pickle.dump(prompt_seq2idx, f)
        
    path = os.path.join(output_dir, "seq2idx_mapping", f'gen_seq2idx.pkl')
    with open(path, 'wb') as f:
        pickle.dump(gen_seq2idx, f)
    
    return prompt_seq2idx, gen_seq2idx


EM_THRESHOLD=0.9

def add_wandb_summary_metrics(accelerator, num_train_steps_per_epoch):
    exact_match_values = EXACT_MATCH_VALUES

    # Find the first crossing timestep and maximum value
    first_crossing_timestep = None
    first_crossing_epoch = None
    max_value = float('-inf')
    max_value_timestep = None
    max_em_value_epoch = None

    for timestep, value in exact_match_values:
        if value > EM_THRESHOLD and first_crossing_timestep is None:
            first_crossing_timestep = timestep
        if value > max_value:
            max_value = value
            max_value_timestep = timestep
    
    if first_crossing_timestep is not None:
        first_crossing_epoch = first_crossing_timestep // num_train_steps_per_epoch

    if max_value_timestep is not None:
        max_em_value_epoch = max_value_timestep // num_train_steps_per_epoch

    accelerator.log({
        "first_crossing_timestep_0.9_EM": first_crossing_timestep,
        "first_crossing_epoch_0.9_EM": first_crossing_epoch,
        "max_exact_match_value": max_value,
        "max_exact_match_timestep": max_value_timestep,
        "max_exact_match_epoch": max_em_value_epoch
    })