# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_model_lm_main.ipynb.

# %% ../nbs/03_model_lm_main.ipynb 3
from __future__ import annotations
import os
from transformers import Trainer, TrainingArguments, AutoConfig
from datasets import DatasetDict,Dataset
import torch
import gc
import math
from functools import partial
import evaluate
from .utils import *
from .text_main_lm import TextDataLMController
from .text_main_lm_streaming import TextDataLMControllerStreaming
from transformers import pipeline

# %% auto 0
__all__ = ['language_model_init', 'compute_lm_accuracy', 'preprocess_lm_logits_for_metrics', 'finetune_lm',
           'extract_hidden_states', 'ModelLMController']

# %% ../nbs/03_model_lm_main.ipynb 5
def language_model_init(model_class, # Model's class object, e.g. AutoModelForMaskedLM
                        cpoint_path=None, # Either model string name on HuggingFace, or the path to model checkpoint. Put `None` to train from scratch
                        config=None, # Model config. If not provided, AutoConfig is used to load config from cpoint_path                        
                        device=None, # Device to train on
                        seed=None, # Random seed
                       ):
    """
    To initialize a language model, either masked or casual
    """
    if device is None: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    if config is None:
        config = AutoConfig.from_pretrained(cpoint_path)

    if seed:
        seed_everything(seed)
    if cpoint_path:
        model = model_class.from_pretrained(cpoint_path,config=config).to(device)
    else:
        print('Initiate a new language model from scratch')
        model = model_class.from_config(config)
    
    print(f'Total parameters: {sum(param.numel() for param in model.parameters())}')
    print(f'Total trainable parameters: {sum(param.numel() for param in model.parameters() if param.requires_grad)}')
    
    return model

# %% ../nbs/03_model_lm_main.ipynb 11
def compute_lm_accuracy(eval_preds, # An EvalPrediction object from HuggingFace 
                        is_mlm, # if this is masked language model, set to `True`. If this is casual language model, set to `False`
                     ):
    
    """    
    Reference: https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm.py#L592C35-L592C35
    """
    metric = evaluate.load("accuracy")
    preds, labels = eval_preds
    # preds have the same shape as the labels, after the argmax(-1) has been calculated
    # by preprocess_logits_for_metrics
    if is_mlm:
        labels = labels.reshape(-1)
        preds = preds.reshape(-1)
        mask = labels != -100
        labels = labels[mask]
        preds = preds[mask]
        return metric.compute(predictions=preds, references=labels)
    else:
        # we need to shift the labels
        labels = labels[:, 1:].reshape(-1)
        preds = preds[:, :-1].reshape(-1)
        return metric.compute(predictions=preds, references=labels)

# %% ../nbs/03_model_lm_main.ipynb 12
def preprocess_lm_logits_for_metrics(logits, labels):
    if isinstance(logits, tuple):
        # Depending on the model and config, logits may contain extra tensors,
        # like past_key_values, but logits always come first
        logits = logits[0]
    return logits.argmax(dim=-1) 

# %% ../nbs/03_model_lm_main.ipynb 13
def finetune_lm(lr, # Learning rate
                bs, # Batch size
                wd, # Weight decay
                epochs, # Number of epochs
                ddict, # The HuggingFace datasetdict
                tokenizer,# HuggingFace tokenizer
                o_dir = './tmp_weights', # Directory to save weights
                save_checkpoint=False, # Whether to save weights (checkpoints) to o_dir
                model=None, # NLP model
                model_init=None, # A function to initialize model
                data_collator=None, # HuggingFace data collator
                compute_metrics=None, # A function to compute metric, default to `compute_lm_accuracy`
                grad_accum_steps=2, # The batch at each step will be divided by this integer and gradient will be accumulated over gradient_accumulation_steps steps.
                lr_scheduler_type='cosine',  # The scheduler type to use. Including: linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup
                warmup_ratio=0.1, # The warmup ratio for some lr scheduler
                no_valid=False, # Whether there is a validation set or not
                seed=None, # Random seed
                report_to='none', # The list of integrations to report the results and logs to. Supported platforms are "azure_ml", "comet_ml", "mlflow", "neptune", "tensorboard","clearml" and "wandb". Use "all" to report to all integrations installed, "none" for no integrations.
                trainer_class=None, # You can include the class name of your custom trainer here
            ):
    "The main model training/finetuning function"
    torch.cuda.empty_cache()
    gc.collect()
    
    if seed:
        seed_everything(seed)
        
    training_args = TrainingArguments(o_dir, 
                                     learning_rate=lr, 
                                     warmup_ratio=warmup_ratio,
                                     lr_scheduler_type=lr_scheduler_type, 
                                     fp16=True,
                                     do_train=True,
                                     do_eval= not no_valid,
                                     evaluation_strategy="no" if no_valid else "epoch", 
                                     save_strategy="epoch" if save_checkpoint else 'no',
                                     overwrite_output_dir=True,
                                     gradient_accumulation_steps=grad_accum_steps,
                                     per_device_train_batch_size=bs, 
                                     per_device_eval_batch_size=bs,
                                     num_train_epochs=epochs, weight_decay=wd,
                                     report_to=report_to,
                                     logging_dir=os.path.join(o_dir, 'log') if report_to!='none' else None,
                                     logging_steps = len(ddict["train"]) // bs,
                                     )

    # instantiate trainer
    trainer_class = Trainer if trainer_class is None else trainer_class
    trainer = trainer_class(
        model=model,
        model_init=model_init if model is None else None,
        args=training_args,
        train_dataset=ddict['train'],#.shard(200, 0)
        eval_dataset=ddict['validation'] if not no_valid else None,
        data_collator=data_collator,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics if not no_valid else None,
        preprocess_logits_for_metrics=preprocess_lm_logits_for_metrics if not no_valid else None
    )
    
    trainer.train()
    return trainer

# %% ../nbs/03_model_lm_main.ipynb 15
def extract_hidden_states(batch,
                          model=None, # NLP model
                          model_input_names=['input_ids', 'token_type_ids', 'attention_mask'], # Model required inputs, from tokenizer.model_input_names
                          data_collator=None, # HuggingFace data collator
                          state_name='last_hidden_state', # Name of the state to extract
                          state_idx=0, # The index (or indices) of the state to extract
                          device = None, # device that the model is trained on
                          ):
    state_idx = val2iterable(state_idx)
    
    if data_collator is not None:    
# --- Convert from  
# {'input_ids': [tensor([    0, 10444,   244, 14585,   125,  2948,  5925,   368,     2]), 
#                tensor([    0, 16098,  2913,   244,   135,   198, 34629,  6356,     2])]
# 'attention_mask': [tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]), 
#                    tensor([1, 1, 1, 1, 1, 1, 1, 1, 1])]
#                    }
# --- to
# [{'input_ids': tensor([    0, 10444,   244, 14585,   125,  2948,  5925,   368,     2]),
#   'attention_mask': tensor([1, 1, 1, 1, 1, 1, 1, 1, 1])},
#  {'input_ids': tensor([    0, 16098,  2913,   244,   135,   198, 34629,  6356,     2]),
#   'attention_mask': tensor([1, 1, 1, 1, 1, 1, 1, 1, 1])}]

        # remove string text, due to transformer new version       
        collator_inp = []
        ks = [k for k in batch.keys() if k in model_input_names]
        vs = [batch[k] for k in ks]
        for pair in zip(*vs):
            collator_inp.append({k:v for k,v in zip(ks,pair)})
        
        batch = data_collator(collator_inp)
    
    inputs = {k:v.to(device) for k,v in batch.items()
              if k in model_input_names}
            
    # switch to eval mode for evaluation
    if model.training:
        model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        states = outputs[state_name]
        for i in state_idx:
            if isinstance(states,tuple):
                states = states[i]
            else:
                states = states[:,i]
    # Switch back to train mode
    if not model.training:
        model.train()

    return {state_name:states.cpu().numpy()}

# %% ../nbs/03_model_lm_main.ipynb 16
class ModelLMController():
    def __init__(self,
                 model, # NLP language model
                 data_store=None, # a TextDataLMController/TextDataLMControllerStreaming object
                 seed=None, # Random seed
                ):
        self.model = model
        self.data_store = data_store
        self.seed = seed
        
    def fit(self,
            epochs, # Number of epochs
            learning_rate, # Learning rate
            ddict=None, # DatasetDict to fit (will override data_store)
            compute_metrics=None, # A function to compute metric, default to `compute_lm_accuracy`
            batch_size=16, # Batch size
            weight_decay=0.01, # Weight decay
            lr_scheduler_type='cosine', # The scheduler type to use. Including: linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup
            warmup_ratio=0.1, # The warmup ratio for some lr scheduler
            o_dir = './tmp_weights', # Directory to save weights
            save_checkpoint=False, # Whether to save weights (checkpoints) to o_dir
            hf_report_to='none', # The list of HuggingFace-allowed integrations to report the results and logs to
            grad_accum_steps=2, # Gradient will be accumulated over gradient_accumulation_steps steps.
            tokenizer=None, # Tokenizer (to override one in ```data_store```)
            data_collator=None, # Data Collator (to override one in ```data_store```)
            is_mlm=None, # Whether this is masked LM or casual LM
            trainer_class=None, # You can include the class name of your custom trainer here
           ):
        
        if tokenizer is None: tokenizer=check_and_get_attribute(self.data_store,'tokenizer')
        if data_collator is None: data_collator=check_and_get_attribute(self.data_store,'data_collator')
        if ddict is None: ddict = check_and_get_attribute(self.data_store,'main_ddict')
        if is_mlm is None: is_mlm = check_and_get_attribute(self.data_store,'is_mlm')
            
        if compute_metrics is None:
            compute_metrics=partial(compute_lm_accuracy,is_mlm=is_mlm)
        
        if len(set(ddict.keys()) & set(['train','training']))==0:
            raise ValueError("Missing the following key for DatasetDict: train/training")
        no_valid = len(set(ddict.keys()) & set(['validation','val','valid']))==0
        
        trainer = finetune_lm(learning_rate,
                              batch_size,
                              weight_decay,
                              epochs,
                              ddict,
                              tokenizer,
                              o_dir,
                              save_checkpoint=save_checkpoint,
                              model=self.model,
                              data_collator=data_collator,
                              compute_metrics=compute_metrics,
                              grad_accum_steps=grad_accum_steps,
                              lr_scheduler_type=lr_scheduler_type,
                              warmup_ratio=warmup_ratio,
                              no_valid=no_valid,
                              seed=self.seed,
                              trainer_class=trainer_class,
                              report_to=hf_report_to)
        
        if not no_valid:
            eval_results = trainer.evaluate()
            try:
                perplexity = math.exp(eval_results["eval_loss"])
                print(f'Perplexity on validation set: {perplexity:.3f}')
            except OverflowError:
                perplexity = float("inf")
                print(f'Perplexity on validation set: {perplexity}')
            eval_results['perplexity'] = perplexity
        
        self.trainer = trainer
        
    def predict_raw_text(self,
                         content:dict|list|str, # Either a single sentence, list of sentence or a dictionary where keys are metadata, values are list
                         print_result=True, # To whether print the result in readable format, or get the result returned
                         **kwargs, # keyword arguments for HuggingFace's text-generation (for clm) or fill-mask (for mlm)
                        ):
        # Example of kwargs for text-generation:
        # https://huggingface.co/docs/transformers/v4.33.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
        
        if not isinstance(self.data_store,(TextDataLMController,TextDataLMControllerStreaming)) or not self.data_store._processed_call:
            raise ValueError('This functionality needs a TextDataController object which has processed some training data')
        test_dset = self.data_store.prepare_test_dataset_from_raws(content)
        result = self.predict_ddict(test_dset,**kwargs)
        if print_result:
            is_mlm = check_and_get_attribute(self.data_store,'is_mlm')
            for preds in result:
                for pred in preds:
                    if is_mlm:
                        print(f"Score: {pred['score']:.3f} >>> {pred['sequence']}")
                    else:
                        print(f">>> {pred['generated_text']}")
                print('-'*20)
        else:
            return result
    
    def predict_ddict(self,
                      dset:Dataset, # A processed and tokenized Dataset
                      **kwargs, # keyword arguments for HuggingFace's text-generation (for clm) or fill-mask (for mlm)
                     ):
        is_mlm = check_and_get_attribute(self.data_store,'is_mlm')
        tokenizer=check_and_get_attribute(self.data_store,'tokenizer')
        main_text=check_and_get_attribute(self.data_store,'main_text')
        _task = 'fill-mask' if is_mlm else 'text-generation'
        pipeline_obj = pipeline(_task,model=self.model,tokenizer=tokenizer,device=self.model.device)
        str_list = dset[main_text]
        if _task=='fill-mask':
            all_tfms = self.data_store.content_tfms 
            all_tfms = partial(func_all,functions=all_tfms) if len(all_tfms) else lambda x: x
            mask_str = all_tfms(tokenizer.mask_token)
            str_list = [str(s).replace(mask_str,tokenizer.mask_token) for s in str_list]
        return [pipeline_obj(s,**kwargs) for s in str_list]
    
    def get_hidden_states_from_raw_text(self,
                                        content:dict|list|str, # Either a single sentence, list of sentence or a dictionary where keys are metadata, values are list
                                        state_name, # Name of the (hidden) state to extract
                                        state_idx=0, # The index (or indices) of the state to extract. For `hidden_states`, accept multiple values
                                       ):
        if not isinstance(self.data_store,(TextDataLMController,TextDataLMControllerStreaming)) or not self.data_store._processed_call:
            raise ValueError('This functionality needs a TextDataController object which has processed some training data')
        dset = self.data_store.prepare_test_dataset_from_raws(content,do_tokenize=True)
        return self.get_hidden_states(dset,
                                      state_name=state_name,
                                      state_idx=state_idx,
                                      batch_size=1
                                     )

        
    def get_hidden_states_from_raw_dset(self,
                                        dset: Dataset, # A raw HuggingFace dataset
                                        state_name, # Name of the (hidden) state to extract
                                        state_idx=0, # The index (or indices) of the state to extract. For `hidden_states`, accept multiple values
                                        batch_size=16, # GPU batch size
                                       ):
        if not isinstance(self.data_store,(TextDataLMController,TextDataLMControllerStreaming)) or not self.data_store._processed_call:
            raise ValueError('This functionality needs a TextDataController object which has processed some training data')
        dset = self.data_store.prepare_test_dataset(dset,do_tokenize=True)
        return self.get_hidden_states(dset,
                                      state_name=state_name,
                                      state_idx=state_idx,
                                      batch_size=batch_size
                                     )
        
        
    def get_hidden_states(self,
                          ddict:DatasetDict|Dataset=None, # A processed and tokenized DatasetDict/Dataset (will override one in ```data_store```)
                          ds_type='test', # The split of DatasetDict to predict
                          state_name='last_hidden_state', # Name of the (hidden) state to extract
                          state_idx=0, # The index (or indices) of the state to extract. For `hidden_states`, accept multiple values
                          batch_size=16, # GPU batch size
                         ):
        
        tokenizer=check_and_get_attribute(self.data_store,'tokenizer')
        if ddict is None: ddict = check_and_get_attribute(self.data_store,'main_ddict')
        if isinstance(ddict,DatasetDict):
            if ds_type not in ddict.keys():
                raise ValueError(f'{ds_type} is not in the given DatasetDict')
            ddict = ddict[ds_type]
        
        ddict.set_format("torch",
                        columns=tokenizer.model_input_names)
        
        results = ddict.map(partial(extract_hidden_states,
                                 model=self.model,
                                 model_input_names = tokenizer.model_input_names,
                                 state_name=state_name,
                                 state_idx=state_idx,
                                 device = self.model.device),
                           batched=True,
                           batch_size=batch_size
                          )
        results.set_format('numpy')
        return results

