-
Notifications
You must be signed in to change notification settings - Fork 2
/
test.py
74 lines (60 loc) · 2.48 KB
/
test.py
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
from argparse import ArgumentParser
from collections import defaultdict
from pathlib import Path
from typing import List, Mapping
import numpy as np
from sklearn.metrics import classification_report
from sklearn.model_selection import KFold
from process_av_data import AVDataset
from utils import MODEL_MAPPING
def main() -> None:
parser = ArgumentParser()
parser.add_argument("--input", type=Path)
parser.add_argument("--model", "-m", choices=["av2e", "emotion", "av", "clark"], type=str)
args = parser.parse_args()
av_dataset = AVDataset.from_json(args.input)
kf = KFold(n_splits=10, shuffle=True, random_state=42)
if args.model in ["av2e", "emotion", "clark"]:
target = "emotions"
all_runs: List[float] = []
elif args.model == "av":
target = "appraisal_variables"
all_runs_vars: Mapping[str, List[float]] = defaultdict(list)
for train_i, test_i in kf.split(av_dataset.text_data):
text_train, text_test = av_dataset.text_data[train_i], av_dataset.text_data[test_i]
var_train, var_test = (
av_dataset.appraisal_variable_data[train_i],
av_dataset.appraisal_variable_data[test_i],
)
em_train, em_test = av_dataset.emotion_data[train_i], av_dataset.emotion_data[test_i]
model = MODEL_MAPPING[args.model].initialize()
model.fit(text_train, var_train, em_train)
source_test = var_test if args.model == "av2e" else text_test
if target == "appraisal_variables":
pred = model.predict(source_test)
for i, var in enumerate(model.variables_encoder.classes_):
report = classification_report(
var_test[:, 2, i],
pred[:, i],
digits=3,
output_dict=True,
zero_division=1.0,
)
all_runs_vars[var].append(report["accuracy"])
elif target == "emotions":
pred = model.predict(source_test)
report = classification_report(
em_test[:, 2],
pred,
digits=3,
output_dict=True,
zero_division=1.0,
)
all_runs.append(report["accuracy"])
if target == "appraisal_variables":
for var in model.variables_encoder.classes_:
print(f"{var}: {np.mean(all_runs_vars[var])}")
elif target == "emotions":
print(np.mean(all_runs))
if __name__ == "__main__":
main()