74 lines
1.9 KiB
Python
74 lines
1.9 KiB
Python
|
import pickle
|
||
|
|
from config import MODEL
|
||
|
|
from scipy.special import softmax
|
||
|
|
|
||
|
|
with open('train_dataset.pickle','rb') as f_p:
|
||
|
|
train_dataset = pickle.load(f_p)
|
||
|
|
|
||
|
|
with open('eval_dataset.pickle','rb') as f_p:
|
||
|
|
eval_dataset = pickle.load(f_p)
|
||
|
|
|
||
|
|
with open('test_dataset.pickle','rb') as f_p:
|
||
|
|
test_dataset = pickle.load(f_p)
|
||
|
|
|
||
|
|
|
||
|
|
from transformers import AutoModelForSequenceClassification
|
||
|
|
|
||
|
|
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)
|
||
|
|
|
||
|
|
from transformers import TrainingArguments
|
||
|
|
|
||
|
|
|
||
|
|
training_args = TrainingArguments("test_trainer",
|
||
|
|
per_device_train_batch_size=4,
|
||
|
|
per_device_eval_batch_size=4,
|
||
|
|
evaluation_strategy='steps',
|
||
|
|
eval_steps=2_000,
|
||
|
|
gradient_accumulation_steps=10,
|
||
|
|
learning_rate=2e-6,
|
||
|
|
warmup_steps=4_000,
|
||
|
|
num_train_epochs=10,
|
||
|
|
load_best_model_at_end=True)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from datasets import load_metric
|
||
|
|
|
||
|
|
metric = load_metric("f1")
|
||
|
|
|
||
|
|
def compute_metrics(eval_pred):
|
||
|
|
logits, labels = eval_pred
|
||
|
|
predictions = np.argmax(logits, axis=-1)
|
||
|
|
return metric.compute(predictions=predictions, references=labels)
|
||
|
|
|
||
|
|
|
||
|
|
from transformers import Trainer
|
||
|
|
|
||
|
|
trainer = Trainer(
|
||
|
|
model=model,
|
||
|
|
args=training_args,
|
||
|
|
train_dataset=train_dataset,
|
||
|
|
eval_dataset=eval_dataset,
|
||
|
|
compute_metrics=compute_metrics,
|
||
|
|
)
|
||
|
|
|
||
|
|
trainer.train()
|
||
|
|
trainer.save_model('./roberta_retrained')
|
||
|
|
trainer.evaluate()
|
||
|
|
|
||
|
|
|
||
|
|
eval_predictions = trainer.predict(eval_dataset).predictions
|
||
|
|
|
||
|
|
with open('../dev-0/out.tsv', 'w') as f_out:
|
||
|
|
for pred in eval_predictions:
|
||
|
|
pred = softmax(pred)[1]
|
||
|
|
f_out.write(str(pred) + '\n')
|
||
|
|
|
||
|
|
test_predictions = trainer.predict(test_dataset).predictions
|
||
|
|
with open('../test-A/out.tsv', 'w') as f_out:
|
||
|
|
for pred in test_predictions:
|
||
|
|
pred = softmax(pred)[1]
|
||
|
|
f_out.write(str(pred) + '\n')
|