108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
import pandas as pd
|
|
import os
|
|
import sys
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
|
|
import torch
|
|
import csv
|
|
|
|
IN_FILE_NAME = "in.tsv.xz"
|
|
OUT_FILE_NAME = "out.tsv"
|
|
TRAIN_PATH = "train"
|
|
EXP_FILE_NAME = "expected.tsv"
|
|
FILE_SEP = "\t"
|
|
IN_HEADER_FILE_NAME = "in-header.tsv"
|
|
OUT_HEADER_FILE_NAME = "out-header.tsv"
|
|
# PT_MODEL_NAME = "bert-base-cased"
|
|
PT_MODEL_NAME = "roberta-base"
|
|
|
|
|
|
class CustomDataset(torch.utils.data.Dataset):
|
|
def __init__(self, encodings, labels):
|
|
self.encodings = encodings
|
|
self.labels = labels
|
|
|
|
def __getitem__(self, idx):
|
|
item = {key: torch.tensor(val[idx])
|
|
for key, val in self.encodings.items()}
|
|
item['labels'] = torch.tensor(self.labels[idx])
|
|
return item
|
|
|
|
def __len__(self):
|
|
return len(self.labels)
|
|
|
|
|
|
def main(dirnames):
|
|
check_path(IN_HEADER_FILE_NAME)
|
|
in_cols = (pd.read_csv(IN_HEADER_FILE_NAME, sep=FILE_SEP)).columns
|
|
check_path(OUT_HEADER_FILE_NAME)
|
|
out_cols = (pd.read_csv(OUT_HEADER_FILE_NAME, sep=FILE_SEP)).columns
|
|
|
|
print("Reading train data...")
|
|
train_set_features = get_tsv_data(os.path.join(
|
|
TRAIN_PATH, IN_FILE_NAME), names=in_cols)
|
|
train_set_labels = get_tsv_data(os.path.join(
|
|
TRAIN_PATH, EXP_FILE_NAME), names=out_cols, compression=None)
|
|
|
|
print("Reading input data...")
|
|
in_sets = []
|
|
for d in dirnames:
|
|
print(f"\tReading dir: {d}...")
|
|
in_sets.append(get_tsv_data(
|
|
os.path.join(d, IN_FILE_NAME), names=in_cols))
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(PT_MODEL_NAME)
|
|
model = AutoModelForSequenceClassification.from_pretrained(
|
|
PT_MODEL_NAME, num_labels=2)
|
|
train_set_enc = tokenizer(
|
|
[t for t in train_set_features[in_cols].agg(' '.join, axis=1)], truncation=True, padding=True)
|
|
dataset = CustomDataset(
|
|
train_set_enc, [int(t) for t in train_set_labels[out_cols[0]]])
|
|
|
|
trainer = Trainer(
|
|
model=model,
|
|
args=TrainingArguments(
|
|
output_dir='./res',
|
|
num_train_epochs=5,
|
|
per_device_train_batch_size=16,
|
|
per_device_eval_batch_size=16
|
|
),
|
|
train_dataset=dataset
|
|
)
|
|
|
|
trainer.train()
|
|
|
|
for i in range(len(in_sets)):
|
|
p = os.path.join(dirnames[i], IN_FILE_NAME)
|
|
with open(p) as f:
|
|
print(
|
|
f"\tPredicting for: {p}...")
|
|
X = [t for t in in_sets[i][in_cols].agg(' '.join, axis=1)]
|
|
out_file_path = os.path.join(dirnames[i], OUT_FILE_NAME)
|
|
f.write('\n'.join(trainer.predict(X)))
|
|
print(f"Saved predictions to file: {out_file_path}")
|
|
|
|
|
|
def get_tsv_data(filename: str, names, compression="infer"):
|
|
check_path(filename)
|
|
return pd.read_csv(
|
|
filename,
|
|
sep=FILE_SEP,
|
|
compression=compression,
|
|
error_bad_lines=False,
|
|
quoting=csv.QUOTE_NONE,
|
|
header=None,
|
|
names=names,
|
|
dtype=str
|
|
)
|
|
|
|
|
|
def check_path(filename: str):
|
|
if not os.path.exists(filename):
|
|
raise Exception(f"Path {filename} does not exist!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
raise Exception("Name of working dir not specified!")
|
|
main(sys.argv[1:])
|