test_multiClass.py

import argparse

from sklearn.model_selection import KFold, train_test_split

from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification, logging

from utils import load_binary_training_data, datasetify, ProbTrainer
from utils import find_best_trial_hyperparam

from torch import cuda

from utils import load_multiClass_training_data, datasetify, ProbTrainer, getPredictionForTrainTest_binary

logging.set_verbosity_warning()

parser = argparse.ArgumentParser()
parser.add_argument('--in_train', dest='in_train', type=str,
                    help='Add filename for training data')
parser.add_argument('--feature', dest='FEATURE', nargs='?', const="",
                    help='Column name of feature, default is "" which leads to a creation of "title_abstract-feature" ')
parser.add_argument('--label', dest='LABEL', nargs='?', const="technologies",
                    help='''List of column names with labels or a string hinting to the list. 
                    If cdr_map_training.csv is used for training, 
                    strings like "technologies", "methodologies", "context", "assessment" result in a list of 
                    meaningful column names. Default is "technologies".
                    ''')
parser.add_argument('--in_model', dest='in_model', type=str,
                    help='Add filename for untrained model')
parser.add_argument('--balanced', dest='balanced', action='store_true',
                    help='Flag if training should use balanced training data')
parser.add_argument('--no_balanced', dest='balanced', action='store_false',
                    help='Flag if training should use no balanced training data')
parser.add_argument('--optimal_hyperparameter', dest='optimal_hyperparameter', action='store_true',
                    help='Flag if hyperparameter should be optimized')
parser.add_argument('--no_optimal_hyperparameter', dest='optimal_hyperparameter', action='store_false',
                    help='Flag if hyperparameter should not be optimized')
parser.add_argument('--folds', dest='folds', type=int, nargs='?', const=4, default=4,
                    help='Number of folds for cross-validation. The complete dataset is used.'
                    )
parser.add_argument('--out_hyperparam', dest='out_hyperparam', type=str,
                    help='Add filename for storage of test runs during hyperaram search.')
parser.add_argument('--out_train', dest='out_train', type=str,
                    help='Add filename for predictions on training data')
parser.add_argument('--out_test', dest='out_test', type=str,
                    help='Add filename for predictions on test dataset')
parser.add_argument('--out_tmp_files', dest='out_tmp_files', type=str,
                    help='Add directory for storage of large tmp-files used for hyperparam-tuning.')


if __name__ == "__main__":
    args = parser.parse_args()
    LABEL = args.LABEL
    FEATURE = args.FEATURE
    FOLDS = args.folds

    print('hyperparameter optimization:', args.optimal_hyperparameter)
    print("label:", LABEL)
    print("feature:", FEATURE)
    print('folds:', FOLDS)
    
    all_train_df, all_classes, FEATURE = load_multiClass_training_data(args.in_train, LABEL=LABEL, FEATURE=FEATURE)
    
    print("label:", all_classes)
    print("feature:", FEATURE)
    
    tokenizer = AutoTokenizer.from_pretrained(args.in_model)
    model = AutoModelForSequenceClassification.from_pretrained(args.in_model, num_labels=2)
    if cuda.is_available():
        model.cuda()

    all_test_runs = []
    all_train_runs = []

    kf = KFold(n_splits=FOLDS, shuffle=True, random_state=43581)
    count = 0
    for train_index, test_index in kf.split(all_train_df):

        print(f"Round {count} of {FOLDS}")
        
        train, test = all_train_df.iloc[train_index], all_train_df.iloc[test_index]
        
        if args.optimal_hyperparameter:
            print("Hyperparameter search...")

            # eval is used to find optimal hyperparameter
            train_small, eval = train_test_split(train, test_size=0.1, random_state=43581)
            best_trial = find_best_trial_hyperparam(train_df=train_small, 
                                            test_df=eval, 
                                            model=model,
                                            tokenizer=tokenizer,
                                            balanced=args.balanced,
                                            model_tmp_dir=f"{args.out_tmp_files}/tmp_trainer", 
                                            ray_log_dir=f"{args.out_tmp_files}/raytune",
                                            FEATURE=FEATURE,
                                            label=all_classes)

            best_trial_params = best_trial.hyperparameters
            with open(args.out_hyperparam, "a") as f:
                f.write(f'run: {count}\n')
                f.write(f'{best_trial_params}\n\n')

        print("\Training...")
        trainer = ProbTrainer(model=model,
                              train_dataset=datasetify(train[FEATURE], 
                                                       tokenizer, 
                                                       train[all_classes].values),
                              balanced=args.balanced,
                              )
        if args.optimal_hyperparameter:
            trainer.train(trial=best_trial_params)
        else:
            trainer.train()

        print("\nEvaluating...")
        test_one_run = getPredictionForTrainTest_multiclass(test, trainer, allClasses, FEATURE, count)
        train_one_run = getPredictionForTrainTest_multiclass(train, trainer, allClasses, FEATURE, count)

        all_test_runs.append(test_one_run)
        all_train_runs.append(train_one_run)

        count += 1

    # store test result
    all_test_runs_df = pd.concat(all_test_runs)
    all_train_runs_df = pd.concat(all_train_runs)

    all_test_runs_df.to_csv(args.out_test, index=False)
    all_train_runs_df.to_csv(args.out_train, index=False)

    print("Predictions for test run on test sets saved to", args.out_test)
    print("Predictions for test run on train sets saved to", args.out_train)