"""
Original Copyright 2022 Loureiro  Daniel, Camacho-collados Jose, et al. Released without any restrictions with scoring code under Apache License 2.0.
Modifications Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
"""

import numpy as np
import evaluate

from transformers import pipeline, AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments, Trainer
from datasets import load_dataset, concatenate_datasets


def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels, average='macro')


dataset = load_dataset('tweet_eval', 'sentiment')


MODEL_NAME = 'cardiffnlp/twitter-roberta-base-sep2022'  # change to desired model from the hub
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# augment train set with test set, for downstream apps only - DO NOT EVALUATE ON TEST
# tokenized_datasets['train+test'] = concatenate_datasets([tokenized_datasets['train'],
#                                                          tokenized_datasets['test']])
id2label = {"0": "negative", "1": "neutral", "2": "positive"}
label2id = {"negative": 0, "neutral": 1, "positive": 2}
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3, id2label=id2label, label2id=label2id)

training_args = TrainingArguments(
    do_eval=True,
    evaluation_strategy='epoch',
    output_dir='test_trainer',
    logging_dir='test_trainer',
    logging_strategy='epoch',
    save_strategy='epoch',
    num_train_epochs=10,
    learning_rate=1e-05,
    # pytorch 2.0.0 specific args
    torch_compile=True,
    bf16=True,
    optim='adamw_torch_fused',
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    load_best_model_at_end=True,
    metric_for_best_model='recall',
)

metric = evaluate.load('recall')  # default metric for sentiment dataset is recall (macro)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
    compute_metrics=compute_metrics,
)

trainer.train()
trainer.create_model_card()

pipeline = pipeline('sentiment-analysis', model=model, tokenizer=tokenizer)
pipeline.save_pretrained('/workspace/aws-deeplearning-labs/workshop/twitter_lm/scripts/saved_model')
# res = trainer.evaluate(tokenized_datasets['test'])
# print(res)