cdr_map_classifier/train_binaryClass.py at master · sarahlbe/cdr_map_classifier · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
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

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="0 relevant-relevant",
                    help='Column name of label, default is "0 relevant-relevant"')
parser.add_argument('--in_model', dest='in_model', type=str,
                    help='Add filename for untrained model')
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('--out_model', dest='out_model', type=str,
                    help='File name of trained model'
                    )
parser.add_argument('--out_train', dest='out_train', type=str,
                    help='Add filename for prediction on training data')
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_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

    print('hyperparameter optimization:', args.optimal_hyperparameter)
    print("label:", LABEL)
    print("feature:", FEATURE)

    train_df, LABEL, FEATURE = load_binary_training_data(args.in_train, LABEL=LABEL, FEATURE=FEATURE)

    print("label:", LABEL)
    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()


    if args.optimal_hyperparameter:
        print("Hyperparameter search...")

        # eval is used to find optimal hyperparameter
        train_small, eval = train_test_split(train_df, 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=False,
                                                model_tmp_dir=f"{args.out_tmp_files}/tmp_trainer",
                                                ray_log_dir=f"{args.out_tmp_files}/raytune",
                                                FEATURE=FEATURE,
                                                label=LABEL)

        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[LABEL].values),
        balanced=False,
        )

    if args.optimal_hyperparameter:
        trainer.train(trial=best_trial_params)
    else:
        trainer.train()

    trainer.save_model(args.out_model)
    print("Model saved to ", args.out_model)