import argparse
import numpy as np

import os
import torch
from transformers import AutoTokenizer, AutoConfig, Trainer, TrainingArguments

from model import RobertaForStsRegression
from dataset import KlueStsWithSentenceMaskDataset
from utils import read_json, seed_everything
from metric import compute_metrics


def main(args):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    config = AutoConfig.from_pretrained(args.model_name_or_path)
    config.num_labels = args.num_labels
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)

    train_file_path = os.path.join(args.data_dir, args.train_filename)
    valid_file_path = os.path.join(args.data_dir, args.valid_filename)

    train_json = read_json(train_file_path)
    valid_json = read_json(valid_file_path)

    train_dataset = KlueStsWithSentenceMaskDataset(train_json, tokenizer, 510)
    valid_dataset = KlueStsWithSentenceMaskDataset(train_json, tokenizer, 510)

    model = RobertaForStsRegression.from_pretrained(
        args.model_name_or_path, config=config
    )
    model.to(device)

    training_args = TrainingArguments(
        output_dir=args.model_dir,
        save_total_limit=args.save_total_limit,
        save_steps=args.save_steps,
        num_train_epochs=args.num_train_epochs,
        learning_rate=args.learning_rate,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=64,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        weight_decay=args.weight_decay,
        logging_dir="./logs",
        logging_steps=args.save_steps,
        evaluation_strategy=args.evaluation_strategy,
        metric_for_best_model="pearsonr",
        fp16=True,
        fp16_opt_level="O1",
        eval_steps=args.save_steps,
        load_best_model_at_end=True,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        compute_metrics=compute_metrics,
    )
    trainer.train()
    model.save_pretrained(args.model_dir)
    tokenizer.save_pretrained(args.model_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # data_arg
    parser.add_argument("--data_dir", type=str, default="./data")
    parser.add_argument("--model_dir", type=str, default="./model")
    parser.add_argument("--output_dir", type=str, default="./output")
    parser.add_argument("--model_name_or_path", type=str, default="klue/roberta-large")
    parser.add_argument(
        "--train_filename", type=str, default="klue-sts-v1.1_train.json"
    )
    parser.add_argument("--valid_filename", type=str, default="klue-sts-v1.1_dev.json")

    # train_arg
    parser.add_argument("--num_labels", type=int, default=1)
    parser.add_argument("--seed", type=int, default=15)
    parser.add_argument("--num_train_epochs", type=int, default=5)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--learning_rate", type=float, default=5e-5)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--weight_decay", type=float, default=0.01)

    # eval_arg
    parser.add_argument("--evaluation_strategy", type=str, default="steps")
    parser.add_argument("--save_steps", type=int, default=250)
    parser.add_argument("--eval_steps", type=int, default=250)
    parser.add_argument("--save_total_limit", type=int, default=2)

    args = parser.parse_args()
    main(args)

[((16, 13, 16, 64), 'transformers.AutoConfig.from_pretrained', 'AutoConfig.from_pretrained', ({(16, 40, 16, 63): 'args.model_name_or_path'}, {}), '(args.model_name_or_path)', False, 'from transformers import AutoTokenizer, AutoConfig, Trainer, TrainingArguments\n'), ((18, 16, 18, 70), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', ({(18, 46, 18, 69): 'args.model_name_or_path'}, {}), '(args.model_name_or_path)', False, 'from transformers import AutoTokenizer, AutoConfig, Trainer, TrainingArguments\n'), ((20, 22, 20, 70), 'os.path.join', 'os.path.join', ({(20, 35, 20, 48): 'args.data_dir', (20, 50, 20, 69): 'args.train_filename'}, {}), '(args.data_dir, args.train_filename)', False, 'import os\n'), ((21, 22, 21, 70), 'os.path.join', 'os.path.join', ({(21, 35, 21, 48): 'args.data_dir', (21, 50, 21, 69): 'args.valid_filename'}, {}), '(args.data_dir, args.valid_filename)', False, 'import os\n'), ((23, 17, 23, 43), 'utils.read_json', 'read_json', ({(23, 27, 23, 42): 'train_file_path'}, {}), '(train_file_path)', False, 'from utils import read_json, seed_everything\n'), ((24, 17, 24, 43), 'utils.read_json', 'read_json', ({(24, 27, 24, 42): 'valid_file_path'}, {}), '(valid_file_path)', False, 'from utils import read_json, seed_everything\n'), ((26, 20, 26, 78), 'dataset.KlueStsWithSentenceMaskDataset', 'KlueStsWithSentenceMaskDataset', ({(26, 51, 26, 61): 'train_json', (26, 63, 26, 72): 'tokenizer', (26, 74, 26, 77): '510'}, {}), '(train_json, tokenizer, 510)', False, 'from dataset import KlueStsWithSentenceMaskDataset\n'), ((27, 20, 27, 78), 'dataset.KlueStsWithSentenceMaskDataset', 'KlueStsWithSentenceMaskDataset', ({(27, 51, 27, 61): 'train_json', (27, 63, 27, 72): 'tokenizer', (27, 74, 27, 77): '510'}, {}), '(train_json, tokenizer, 510)', False, 'from dataset import KlueStsWithSentenceMaskDataset\n'), ((29, 12, 31, 5), 'model.RobertaForStsRegression.from_pretrained', 'RobertaForStsRegression.from_pretrained', (), '', False, 'from model import RobertaForStsRegression\n'), ((34, 20, 52, 5), 'transformers.TrainingArguments', 'TrainingArguments', (), '', False, 'from transformers import AutoTokenizer, AutoConfig, Trainer, TrainingArguments\n'), ((54, 14, 60, 5), 'transformers.Trainer', 'Trainer', (), '', False, 'from transformers import AutoTokenizer, AutoConfig, Trainer, TrainingArguments\n'), ((67, 13, 67, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((15, 38, 15, 63), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n')]