utils.py

import pandas as pd
from ast import literal_eval

from datasets import Dataset
from transformers import AutoTokenizer, TrainingArguments
from transformers import AutoModelForSequenceClassification, Trainer, logging

from torch import tensor
from torch.nn import Sigmoid, Softmax
from transformers.trainer_utils import PredictionOutput
from ray import tune

def load_binary_training_data(filename, LABEL, FEATURE):
    
    if len(LABEL)==0:
        # label from cdr_map_training.csv
        LABEL = '0 relevant-relevant'
    
    if filename.split(".")[-1]=="csv":
        df = pd.read_csv(filename)
        
    df = df.loc[(df[LABEL] == 0) | (df[LABEL] == 1)]
    if len(FEATURE)==0:
        FEATURE = "title_abstract"
        df[FEATURE] = df["title"] + ". " + df["abstract"]
    
    df = df[["doc_id", FEATURE, LABEL]].copy()
    
    # some cleaning
    df = df.loc[(df[LABEL]==0)|(df[LABEL]==1)]
    df = df.dropna(subset=[FEATURE])    
    
    return df, LABEL, FEATURE

def load_multiClass_training_data(filename, LABEL, FEATURE):
       
    if filename.split(".")[-1]=="csv":
        df = pd.read_csv(filename)
    
    if len(FEATURE)==0:
        FEATURE = "title_abstract"
        df[FEATURE] = df["title"] + ". " + df["abstract"]
    
    try:
        LABEL = literal_eval(LABEL)
    except ValueError:
        pass
    
    if isinstance(LABEL, list):
        # label from cdr_map_training.csv
        all_classes = LABEL
    else: 
        all_classes = [c for c in df.columns if LABEL in c]

    # some cleaning
    if "0_relevant-relevant" in df.columns:
        df = df.loc[df["0_relevant-relevant"]==1].copy()
    for c in all_classes:
        df = df.loc[(df[c]==0)|(df[c]==1)].copy()    
    df = df.dropna(subset=[FEATURE])    
    df = df[["doc_id", FEATURE]+all_classes].copy()
    
    
    return df, all_classes, FEATURE



class ProbTrainer(Trainer):
    def __init__(self, balanced, *args, **kwargs):
        self.balanced = balanced
        super().__init__(*args, **kwargs)

    def predict_proba(self, test_dataset: Dataset) -> PredictionOutput:
        logits = self.predict(test_dataset).predictions
        if logits.shape[1] > 2:
            activation = Sigmoid()
        else:
            activation = Softmax()
        return activation(tensor(logits)).numpy()

    # used if balancing class weights
    def __calculate_weights(self):
        labels = self.train_dataset["label"]
        labels = pd.DataFrame(labels)

        weights = []
        for i in labels.columns:
            # weight_i = number_of_not_class/number_of_class
            try:
                weight_i = (len(labels) - sum(labels[i])) / sum(labels[i])
            except ZeroDivisionError:
                print('No samples in class')
                weight_i = (len(labels) - sum(labels[i])) / 1.
            weights.append(weight_i)
        return tensor(weights)

    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        if cuda.is_available():
            labels.cuda()
            logits.cuda()

        if self.balanced:
            # compute custom loss
            loss_fct = nn.CrossEntropyLoss(weight=self.__calculate_weights())
            if cuda.is_available():
                loss_fct.cuda()
            loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels)

        if not self.balanced:
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        return (loss, outputs) if return_outputs else loss



def find_best_trial_hyperparam(train_df, test_df, model, tokenizer, balanced, model_tmp_dir, ray_log_dir, FEATURE, label):
    training_args = TrainingArguments(
        #"test"
        evaluation_strategy="steps",
        eval_steps=1e9,
        disable_tqdm=True,
        output_dir=model_tmp_dir#"/p/tmp/salueck/tmp_trainer"
    )

    def model_init(trial):
        return model

    trainer = ProbTrainer(model=None,
                          train_dataset=datasetify(train_df[FEATURE], tokenizer, train_df[label].values),
                          eval_dataset=datasetify(test_df[FEATURE], tokenizer, test_df[label].values),
                          balanced=balanced,
                          model_init=model_init,
                          args=training_args,
                          )

    def ray_hp_space(trial):
        return {
            "learning_rate": tune.loguniform(1e-6, 1e-4),
            "per_device_train_batch_size": tune.choice([16, 32, 64, 128]),
            "weight_decay": tune.uniform(0, 0.3),
            "warmup_steps": tune.loguniform(1, 500),
            "num_epochs": tune.choice([2, 3, 4, 5])
        }

    best_trial = trainer.hyperparameter_search(
        direction="maximize",
        hp_space=ray_hp_space,
        backend="ray",
        n_trials=20,  # number of trials
        local_dir=ray_log_dir, #"/p/tmp/salueck/ray_log/",
        keep_checkpoints_num=1,
        resources_per_trial={"cpu": 14, "gpu": 1},
    )

    return best_trial

# tokenizing
def datasetify(x, tokenizer, y=None):
    data_dict = {"text": x}
    if y is not None:
        data_dict["label"] = y
    dataset = Dataset.from_dict(data_dict)

    def tokenize_function(examples):
        return tokenizer(examples["text"], padding='longest', truncation=True)

    return dataset.map(tokenize_function, batched=True)

def getPredictionForTrainTest_multiclass(df, trainer, allClasses, FEATURE, count):
    
    prediction = trainer.predict_proba(test_dataset=datasetify(df[FEATURE], tokenizer))
    pred_one_run = pd.DataFrame({'id': df["id"],
                                 'seed': [count for i in range(len(df))]}
                                )
    pred_one_run = pred_one_run.merge(df[["id"] + list(allClasses)], on='id', how='left')
    pred_one_run[['pred_' + c for c in allClasses]] = np.where(prediction >= 0.5, 1, 0)
    pred_one_run[['predProba_' + c for c in allClasses]] = prediction
    pred_one_run = pred_one_run.apply(lambda x: getClassWithHighestProbability(x, allClasses), axis=1)

    return pred_one_run

def getPredictionForTrainTest_binary(df, trainer, LABEL, FEATURE, count):
    
    prediction = trainer.predict_proba(test_dataset=datasetify(df[FEATURE], tokenizer))
    pred_one_run = pd.DataFrame({'id': df["id"],
                                 'seed': [count for i in range(len(df))]}
                                )
    pred_one_run = pred_one_run.merge(df[["id", LABEL]], on='id', how='left')
    pred_one_run["predProba"] = prediction[:, 1]
    pred_one_run["pred"] = np.where(prediction[:, 1] >= 0.5, 1, 0)

    return pred_one_run

def predict_multiClass(model_path, tokenizer, allClasses, overall, df, columns_to_keep):
    path = model_path #args.pretrained_model_technology
    model = AutoModelForSequenceClassification.from_pretrained(path, num_labels=len(allClasses))
    trainer = ProbTrainer(model=model)
    pred = trainer.predict_proba(test_dataset=datasetify(df[FEATURE], tokenizer))

    df[['pred_' + c for c in allClasses]] = np.where(pred >= 0.5, 1, 0)
    df[['predProba_' + c for c in allClasses]] = pred
    df[f'all_pred_{overall}'] = df.apply(collect_pred_class, args=([allClasses]), axis=1)
    df[f'first_pred_{overall}'] = df.apply(get_highest_prediction, args=([allClasses]), axis=1)

    columns_to_keep.extend(['pred_' + c for c in allClasses])
    columns_to_keep.extend(['predProba_' + c for c in allClasses])

    columns_to_keep.extend([f'all_pred_{overall}', f'first_pred_{overall}'])

    return df, columns_to_keep

def collect_pred_class(row, allClasses):
    all_pred_classes = []
    for c in allClasses:
        if row[f'predProba_{c}'] >= 0.5:
            all_pred_classes.append(c)
    if len(all_pred_classes)==0:
        predProbaClasses = ['predProba_' + c for c in allClasses]
        max_val = 0
        for predProbaClass in predProbaClasses:
            if row[predProbaClass] > max_val:
                max_val = row[predProbaClass]
                tech = predProbaClass.split("predProba_")[1]
        all_pred_classes.append(tech)
    return all_pred_classes

def get_highest_prediction(row, allClasses):
    max_proba = 0
    for c in allClasses:
        if row[f'predProba_{c}'] > max_proba:
            max_proba = row[f'predProba_{c}']
            max_class = c

    return max_class