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')