apt/train-models-on-ppdb.py at main · Advancing-Machine-Human-Reasoning-Lab/apt · 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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
"""
COMMAND LINE ARGUMENTS -
1. Dataset path
2. Model family
3. Model name (or path for a saved model)
4. Output dir
5. Additional dataset name(s) comma separated
"""

import os
import sys
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.utils import resample
from simpletransformers.classification import ClassificationModel
from utils import read

if not os.path.exists(sys.argv[4]):
    os.makedirs(sys.argv[4])

dataset = sys.argv[1]
data_path = "/raid/datasets/" + dataset + "/"

if "twitter" in dataset:
    train_df = pd.read_csv(data_path + "train-preprocessed.txt", sep="\t", header=None)
    train = pd.DataFrame(
        {
            "text_a": train_df.iloc[:, 0],
            "text_b": train_df.iloc[:, 1],
            "labels": train_df.iloc[:, 2],
        }
    )
    test_df = pd.read_csv(data_path + "test-preprocessed.txt", sep="\t", header=None)
    test = pd.DataFrame(
        {
            "text_a": test_df.iloc[:, 0],
            "text_b": test_df.iloc[:, 1],
            "labels": test_df.iloc[:, 2],
        }
    )
else:
    train_df = pd.read_csv(data_path + "train-preprocessed.txt", sep=" ||| ", header=None)
    df = pd.DataFrame(
        {
            "text_a": train_df.iloc[:, 0],
            "text_b": train_df.iloc[:, 1],
            "labels": train_df.iloc[:, 2],
        }
    )
    train, test = train_test_split(df, test_size=0.1)

print(train.shape)
epochs = 6 if 'large' in sys.argv[3] else 3
try:
    datasets = sys.argv[5].split(",")
    print(datasets)
    epochs = 10 if 'large' in sys.argv[3] else 5
    new_dataset = read(datasets[0], sep="\t", header=True)
    for dataset in datasets[1:]:
        new_dataset = new_dataset.append(read(dataset, sep="\t", header=True))
    print("new dataset", new_dataset.shape)
    try:
        if sys.argv[6] == "balanced":
            new_dataset_mi = new_dataset[new_dataset.iloc[:, -1] == 1]
            new_dataset_nmi = new_dataset[new_dataset.iloc[:, -1] == 0]
            if new_dataset_mi.shape[0] <= new_dataset_nmi.shape[0]:
                new_dataset = pd.concat(
                    [
                        new_dataset_mi,
                        resample(
                            new_dataset_nmi,
                            replace=True,
                            n_samples=len(new_dataset_mi),  # to match minority class
                            random_state=0,  # reproducible results
                        ),
                    ]
                )
            else:
                new_dataset = pd.concat(
                    [
                        new_dataset_nmi,
                        resample(
                            new_dataset_mi,
                            replace=True,
                            n_samples=len(new_dataset_nmi),  # to match minority class
                            random_state=0,  # reproducible results
                        ),
                    ]
                )
            print("balanced new dataset", new_dataset.shape)
    except:
        pass
    train = pd.concat([train, new_dataset])
except:
    pass
print("final dataset", train.shape)

# shuffle
train = train.sample(frac=1).reset_index(drop=True)
test = test.sample(frac=1).reset_index(drop=True)

model = ClassificationModel(
    sys.argv[2],
    sys.argv[3],
    num_labels=2,
    use_cuda=True,
    cuda_device=0,
    args={
        "output_dir": sys.argv[4],
        "overwrite_output_dir": False,
        "fp16": False,  # uses apex
        "num_train_epochs": epochs,
        "train_batch_size": 88,
        "eval_batch_size": 88,
        "do_lower_case": False,
        "evaluate_during_training": True,
        "evaluate_during_verbose": True,
        "evaluate_during_training_steps": 10000,
        "n_gpu": 1,
        "reprocess_input_data": True,
    },
)

model.train_model(train, eval_df=test)