68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
|
||
|
|
import gzip
|
||
|
|
import logging
|
||
|
|
import lzma
|
||
|
|
from typing import List
|
||
|
|
|
||
|
|
from simpletransformers.classification import ClassificationModel
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
MAPPER_LABERL2ID = {
|
||
|
|
'F': 0,
|
||
|
|
'M': 1,
|
||
|
|
}
|
||
|
|
MAPPER_ID2LABEL = dict([(v, k) for k, v in MAPPER_LABERL2ID.items()])
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def open_file(path, *args):
|
||
|
|
if path.endswith('gz'):
|
||
|
|
fopen = gzip.open
|
||
|
|
elif path.endswith('xz'):
|
||
|
|
fopen = lzma.open
|
||
|
|
else:
|
||
|
|
fopen = open
|
||
|
|
return fopen(path, *args)
|
||
|
|
|
||
|
|
|
||
|
|
def load_test(path: str) -> List[str]:
|
||
|
|
data = []
|
||
|
|
|
||
|
|
logger.info(f'Loading {path}')
|
||
|
|
with open_file(path, 'rt') as f:
|
||
|
|
for line in tqdm(f):
|
||
|
|
line = line.strip()
|
||
|
|
data.append(line)
|
||
|
|
|
||
|
|
return data
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
logging.basicConfig(level=logging.INFO)
|
||
|
|
transformers_logger = logging.getLogger('transformers')
|
||
|
|
#transformers_logger.setLevel(logging.WARNING)
|
||
|
|
|
||
|
|
args = {
|
||
|
|
'train_batch_size': 200,
|
||
|
|
'num_train_epochs': 2,
|
||
|
|
'evaluate_during_training': True,
|
||
|
|
'save_steps': 15000,
|
||
|
|
'evaluate_during_training_steps': 15000,
|
||
|
|
}
|
||
|
|
|
||
|
|
model = ClassificationModel('roberta', 'outputs/best_model',
|
||
|
|
num_labels=2, args=args)
|
||
|
|
output_name = 'model=base'
|
||
|
|
|
||
|
|
for test_name in ['dev-0', 'dev-1', 'test-A']:
|
||
|
|
logger.info(f'START TESTING {test_name}')
|
||
|
|
test_data = load_test(f'data/{test_name}/in.tsv.xz')
|
||
|
|
predictions, raw_outputs = model.predict(test_data)
|
||
|
|
logger.info(f'Saving predictions')
|
||
|
|
with open_file(f'data/{test_name}/out-{output_name}.tsv', 'wt') as w:
|
||
|
|
for prediction in predictions:
|
||
|
|
label_name = MAPPER_ID2LABEL[prediction]
|
||
|
|
w.write(f'{label_name}\n')
|