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

# %% ../nbs/03_model_main.ipynb 3
from __future__ import annotations
import os, sys
from transformers import Trainer, TrainingArguments, AutoConfig
from datasets import DatasetDict,Dataset
from pathlib import Path
import torch
import gc
from sklearn.metrics import accuracy_score
from functools import partial
import pandas as pd
import numpy as np
from .utils import *
from .text_main import TextDataMain

# %% auto 0
__all__ = ['model_init_classification', 'compute_metrics_classification', 'compute_metrics_separate_singleheads',
           'loss_for_classification', 'finetune', 'ModelController']

# %% ../nbs/03_model_main.ipynb 4
def model_init_classification(
                              model_class, # Model's class object, e.g. RobertaHiddenStateConcatForSequenceClassification
                              cpoint_path, # Either model string name on HuggingFace, or the path to model checkpoint
                              output_hidden_states:bool, # To whether output the model hidden states or not. Useful when you try to build a custom classification head 
                              device=None, # Device to train on
                              seed=42, # Random seed
                              body_model=None, # If not none, we use this to initialize model's body. If you only want to load the model checkpoint in cpoint_path, leave this as none
                              model_kwargs={} # Keyword arguments for model (both head and body)
                             ):
    """To initialize a classification model, either from an existing HuggingFace model or custom architecture
    
    Can be used for binary, multi-class single-head, multi-class "two-head", and multi-label clasisifcation
    """
    if device is None: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    config = AutoConfig.from_pretrained(
        cpoint_path,
        output_hidden_states=output_hidden_states,
    )
    
    seed_everything(seed)
    if body_model is not None:
        model = model_class(config=config,**model_kwargs)
        layers = list(model.children())
        print('Loading body weights. This assumes the body is the very first first-layer block of your custom architecture')
        body_name, _ = next(iter(model.named_children()))
        setattr(model, body_name, body_model)
        model = model.to(device)
        
    else:
        model = model_class.from_pretrained(cpoint_path,config=config,**model_kwargs).to(device)
    return model

# %% ../nbs/03_model_main.ipynb 6
def compute_metrics_classification(pred, # An EvalPrediction object from HuggingFace (which is a named tuple with ```predictions``` and ```label_ids``` attributes)
                                   metric_funcs=[], # A list of metric functions to evaluate
                                   head_sizes=[], # Class size for each head,
                                   label_names=[], # Names of the label (dependent variable) columns
                                   is_multilabel=False, # Whether this is a multilabel classification
                                   multilabel_threshold=0.5 # Threshold for multilabel (>= threshold is positive)
                                  ):
    """
    Return a dictionary of metric name and its values. Can handle both multiclass and multilabel    
    
    Reference: https://github.com/huggingface/transformers/blob/dbc12269ed5546b2da9236b9f1078b95b6a4d3d5/src/transformers/trainer_utils.py#LL100C22-L100C22
    """
    assert len(head_sizes)==len(label_names)
    labels = pred.label_ids 
    if isinstance(pred.predictions,tuple):
        preds = pred.predictions[0]
    else:
        preds = pred.predictions
    results={}
    metric_funcs = val2iterable(metric_funcs)
    
    for i,(_size,_name) in enumerate(zip(head_sizes,label_names)):
        start= 0 if i==0 else start+head_sizes[i-1]
        end = start + _size
        _pred = preds[:,start:end]
        if is_multilabel:
            # sigmoid and threshold
            _pred = (sigmoid(_pred)>=multilabel_threshold).astype(int)
        else:
            _pred = _pred.argmax(-1)
        _label = labels[:,i] if len(head_sizes)>1 else labels
        for m_func in metric_funcs:
            m_name = callable_name(m_func)
            results[f'{m_name}_{_name}']=m_func(_label,_pred)
    return results

# %% ../nbs/03_model_main.ipynb 8
def compute_metrics_separate_singleheads(pred, # An EvalPrediction object from HuggingFace (which is a named tuple with ```predictions``` and ```label_ids``` attributes)
                              metric_funcs=[], # A list of metric functions to evaluate
                              label_names=[], # Names of the label (dependent variable) columns
                              **kwargs
                             ):
    """
    Return a dictionary of metric name and its values. Can handle both multiclass and multilabel
    """
    # pred: EvalPrediction object 
    # (which is a named tuple with predictions and label_ids attributes)
    labels = pred.label_ids # (bs,number of head separately)
    assert labels.shape[1]==len(label_names)
    
    results={}
    metric_funcs = val2iterable(metric_funcs)
    
    for i in range(len(label_names)):
        _label = labels[:,i]
        _pred = pred.predictions[i].argmax(-1)
        for m_func in metric_funcs:
            m_name = callable_name(m_func)
            results[f'{m_name}_{label_names[i]}']=m_func(_label,_pred)
    
    return results

# %% ../nbs/03_model_main.ipynb 10
def loss_for_classification(logits, # output of the last linear layer, before any softmax/sigmoid. Size: (bs,class_size)
                            labels, # determined by your datasetdict. Size: (bs,number_of_head)
                            is_multilabel=False, # Whether this is a multilabel classification
                            is_multihead=False, # Whether this is a multihead (multi-level) classification
                            head_sizes=[], # class size for each head
                            head_weights=[], # loss weight for each head
                           ):
    """
    The general loss function for classification
    
    - If is_multilabel is ```False``` and is_multihead is ```False```: One-Head Classification, e.g. You predict 1 out of n class
    
    - If is_multilabel is ```False``` and is_multihead is ```True```: Multi-Head Classification, e.g. You predict 1 out of n classes at Level 1, 
    and 1 out of m classes at Level 2
    
    - If is_multilabel is ```True``` and is_multihead is ```False```: One-Head Multi-Label Classification, e.g. You predict x out of n class (x>=1)
    
    - If is_multilabel is ```True``` and is_multihead is ```True```: Not supported!
    
    """
    if is_multilabel and is_multihead: raise ValueError('Multi-Label and Multi-Head problem is not supported')
    head_sizes = val2iterable(head_sizes)
    loss=0
    if not is_multilabel:
        if not is_multihead:
            loss_fct = torch.nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, head_sizes[0]), labels.view(-1))
        else:
            assert len(head_sizes)==len(head_weights),"For MultiHead, make sure len of head_sizes and head_weights equal"
            for i,(_size,_weight) in enumerate(zip(head_sizes,head_weights)):
                start= 0 if i==0 else start+head_sizes[i-1]
                end = start + _size
                loss_fct = torch.nn.CrossEntropyLoss()
                loss = loss + _weight*loss_fct(logits[:,start:end].view(-1,_size),
                                               labels[:,i].view(-1))
    else:
        if not is_multihead:
            loss_fct = torch.nn.BCEWithLogitsLoss()
#             label_1hot = torch.nn.functional.one_hot(labels.view(-1),num_classes=head_sizes[0])
            loss = loss_fct(logits,
                            labels.float())
        else:
            raise ValueError('Multi-Head with Multi-Label classification is not supported. Have you lost your mind?')
#             assert len(head_sizes)==len(head_weights),"For MultiHead, make sure len of head_sizes and head_weights equal"
#             for i,(_size,_weight) in enumerate(zip(head_sizes,head_weights)):
#                 start= 0 if i==0 else start+head_sizes[i-1]
#                 end = start + _size
#                 loss_fct = torch.nn.BCEWithLogitsLoss()
#                 loss = loss + _weight*loss_fct(logits[:,start:end].view(-1,_size),
#                                                torch.nn.functional.one_hot(labels[:,i].view(-1),num_classes=_size).float()
#                                               )
            
    return loss

# %% ../nbs/03_model_main.ipynb 12
def finetune(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, e.g. `compute_metrics_classification`
             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
             no_valid=False, # Whether there is a validation set or not
             seed=42, # 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.
            ):
    "The main model training/finetuning function"
    torch.cuda.empty_cache()
    gc.collect()

    seed_everything(seed)
    training_args = TrainingArguments(o_dir, 
                                learning_rate=lr, 
                                warmup_ratio=0.1 if lr_scheduler_type=='cosine' else 0.0, 
                                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 = Trainer(
        model=model,
        model_init=model_init if model is None else None,
        args=training_args,
        train_dataset=ddict['train'],#.shard(200, 0),    # Only use subset of the dataset for a quick training. Remove shard for full training
        eval_dataset=ddict['validation'] if not no_valid else None,#.shard(100, 0), # Only use subset of the dataset for a quick training. Remove shard for full training
        data_collator=data_collator,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics,
    )
    
    
    trainer.train()
    return trainer

# %% ../nbs/03_model_main.ipynb 14
def _forward_pass_for_predictions(batch,
                                 model=None, # NLP model
                                 topk=1, # Number of labels to return for each head
                                 is_multilabel=False, # Is this a multilabel classification?
                                 multilabel_threshold=0.5, # The threshold for multilabel classification
                                 tokenizer=None, # HuggingFace tokenizer
                                 data_collator=None, # HuggingFace data collator
                                 cols_to_remove=[], # list of keys (columns) to remove from ```batch```
                                 label_names=[], # Names of the label columns
                                 label_sizes=[], # Size of each label
                                 device = None, # device that the model is trained on
                                 is_dhc=False
                                 ):
    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 not in cols_to_remove]
        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 tokenizer.model_input_names}
    
    _f = partial(torch.nn.functional.softmax,dim=1) if not is_multilabel else torch.sigmoid
    
    # switch to eval mode for evaluation
    if model.training:
        model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        outputs_logits = outputs.logits
        outputs_list=[]
        if is_dhc:
            # outputs_logits will be a list of n (typically 2) heads' logits, each has shape (bs,class_size)
            for i in range(len(label_names)):
                outputs_list.append(_f(outputs_logits[i].cpu()))
        else:
            # outputs_logits will have shape (bs,sum of all class sizes). We split into each class
            outputs_logits = outputs_logits.cpu()
            _s=0
            _e=label_sizes[0]
            for i in range(len(label_names)):
                outputs_list.append(_f(outputs_logits[:,_s:_e]))
                _s+=label_sizes[i]
                _e+=label_sizes[i+1] if i+1<len(label_names) else 0
        
        # save prediction and probability
        pred_label_list=[]
        pred_prob_list=[]
        if is_multilabel:
            for i in range(len(label_names)):
                pred_label_list.append(outputs_list[i]>=multilabel_threshold)
                pred_prob_list.append(outputs_list[i])
        else:
            for i in range(len(label_names)):
                _p,_l = torch.topk(outputs_list[i],topk,dim=-1)
                if topk==1:
                    _l,_p = _l[:,0],_p[:,0]
                pred_label_list.append(_l)
                pred_prob_list.append(_p)
    
    # Switch back to train mode
    if not model.training:
        model.train()
        
    results={}
    for i in range(len(label_names)):
        results[f'pred_{label_names[i]}']= pred_label_list[i].numpy()
        results[f'pred_prob_{label_names[i]}']= pred_prob_list[i].numpy()
    return results

# %% ../nbs/03_model_main.ipynb 15
class ModelController():
    def __init__(self,
                 model, # NLP model
                 data_store:TextDataMain=None, # a TextDataMain object
                 metric_funcs=[accuracy_score], # Metric function (can be from Sklearn)
                 seed=42, # Random seed
                ):
        self.model = model
        self.data_store = data_store
        self.metric_funcs = metric_funcs
        self.seed = seed
        
    def fit(self,
            epochs, # Number of epochs
            learning_rate, # Learning rate
            ddict=None, # DatasetDict to fit (will override data_store)
            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
            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
            compute_metrics=None, # A function to compute metric, e.g. `compute_metrics_classification` which utilizes the given ```metric_funcs``` 
            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```)
            label_names=None, # Names of the label (dependent variable) columns (to override one in ```data_store```)
            head_sizes=None, # Class size for each head (to override one in ```model```)
           ):
        
        if tokenizer is None: tokenizer=check_and_get_attribute(self.data_store,'tokenizer')
        if data_collator is None: data_collator=getattr(self.data_store,'data_collator',None)
        if ddict is None: ddict = check_and_get_attribute(self.data_store,'main_ddict')
            
        if label_names is None: label_names=check_and_get_attribute(self.data_store,'label_names')
        label_names = val2iterable(label_names)
        
        if head_sizes is None: head_sizes=check_and_get_attribute(self.model,'head_class_sizes')
        head_sizes = val2iterable(head_sizes)
        
        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']))==0

        _compute_metrics = partial(compute_metrics,
                                   metric_funcs=self.metric_funcs,
                                   head_sizes=head_sizes,
                                   label_names=label_names 
                                  )
        

        trainer = finetune(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,
                           no_valid=no_valid,
                           seed=self.seed,
                           report_to=hf_report_to)
        self.trainer = trainer
        
    def predict_raw_text(self,
                         content:dict|list|str, # Either a single sentence, list of sentence or a dictionary with keys are metadata, values are list
                         batch_size=1, # Batch size. For a small amount of texts, you might want to keep this small
                         is_multilabel=None, # Is this a multilabel classification?
                         multilabel_threshold=0.5, # Threshold for multilabel classification
                         topk=1, # Number of labels to return for each head
                         is_dhc=False # Are outpuf (of model) separate heads?
                        ):
        if not isinstance(self.data_store,TextDataMain) or not self.data_store._main_called:
            raise ValueError('This functionality needs a TextDataMain object which has already processed some training data')
        with HiddenPrints():
            test_ddict = self.data_store.get_test_datasetdict_from_dict(content)
            df_result =  self.predict_ddict(ddict=test_ddict,
                                            ds_type='test',
                                            batch_size=batch_size,
                                            is_multilabel=is_multilabel,
                                            multilabel_threshold=multilabel_threshold,
                                            topk=topk,
                                            is_dhc=is_dhc
                                           )
        return df_result
    
    def predict_ddict(self,
                      ddict=None, # DatasetDict to predict (will override ```data_store```)
                      ds_type='test', # Keys of DatasetDict to predict
                      batch_size=16, # Batch size
                      is_multilabel=None, # Is this a multilabel classification?
                      multilabel_threshold=0.5, # Threshold for multilabel classification
                      topk=1, # Number of labels to return for each head
                      tokenizer=None, # Tokenizer (to override one in ```data_store```)
                      data_collator=None, # Data Collator (to override one in ```data_store```)
                      label_names=None, # Names of the label (dependent variable) columns (to override one in ```data_store```)
                      class_names_predefined=None, # List of names associated with the labels (same index order) (to override one in ```data_store```)
                      device=None, # Device that the model is trained on
                      is_dhc=False # Are outpuf (of model) separate heads?
                     ):
        if device is None: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        if is_multilabel is None: is_multilabel=getattr(self.model,'is_multilabel',False)
        label_lists = class_names_predefined
        
        if tokenizer is None: tokenizer=check_and_get_attribute(self.data_store,'tokenizer')
        if data_collator is None: data_collator=getattr(self.data_store,'data_collator',None)
        if label_names is None: label_names=check_and_get_attribute(self.data_store,'label_names')
        if label_lists is None: label_lists = check_and_get_attribute(self.data_store,'label_lists')
        if not isinstance(label_names,list):
            label_names=[label_names]
        if not isinstance(label_lists[0],list):
            label_lists=[label_lists]    
            
        label_sizes = [len(cs) for cs in label_lists]
        if ddict is None: ddict = check_and_get_attribute(self.data_store,'main_ddict') 
        if ds_type not in ddict.keys():
            raise ValueError(f'{ds_type} is not in the given DatasetDict keys')
        
        ddict.set_format("torch",
                        columns=["input_ids", "attention_mask"])
        
        print_msg('Start making predictions',20)
        preserved_kws=['input_ids', 'token_type_ids', 'attention_mask','label'] # hard-coded
        cols_to_remove = (set(ddict[ds_type].features.keys()) - set(preserved_kws)) | {'text'}
        ddict[ds_type] = ddict[ds_type].map(
            partial(_forward_pass_for_predictions,model=self.model,
                    topk=topk,
                    is_multilabel=is_multilabel,
                    multilabel_threshold=multilabel_threshold,
                    tokenizer=tokenizer,
                    data_collator=data_collator,
                    cols_to_remove=cols_to_remove,
                   label_names=label_names,
                    label_sizes=label_sizes,
                    is_dhc = is_dhc,
                    device=device
                   ), 
            batched=True, batch_size=batch_size)
    
        ddict.set_format("pandas")
        
        df_result = ddict[ds_type][:]
        cols_to_keep = [c for c in df_result.columns.values if c not in preserved_kws[:-1]]
        df_result = df_result.loc[:,cols_to_keep]
        
        # convert pred id to string label
        for i in range(len(label_names)):  
            if not is_multilabel:
                if topk==1:
                    df_result[f'pred_{label_names[i]}'] = df_result[f'pred_{label_names[i]}'].apply(lambda x: label_lists[i][int(x)])
                else:
                    df1 = pd.DataFrame(df_result[f'pred_{label_names[i]}'].to_list(),columns=[f'pred_{label_names[i]}_top{j}' for j in range(1,topk+1)])
                    df1_prob = pd.DataFrame(df_result[f'pred_prob_{label_names[i]}'].to_list(),columns=[f'pred_prob_{label_names[i]}_top{j}' for j in range(1,topk+1)])

                    for j in range(1,topk+1):
                        df1[f'pred_{label_names[i]}_top{j}'] =  df1[f'pred_{label_names[i]}_top{j}'].apply(lambda x: label_lists[i][int(x)])

                    df_result = pd.concat([df_result,df1,df1_prob],axis=1)
            else:
                get_label_str_multilabel = lambda row: ','.join([label_lists[i][int(j)] for j in np.where(row==True)[0]])
                df_result[f'pred_{label_names[i]}_string'] = df_result[f'pred_{label_names[i]}'].apply(get_label_str_multilabel)

        return df_result
