main.py

import torch
from torch import nn

from transformers import AutoTokenizer, AdamW, HfArgumentParser, TrainingArguments, get_linear_schedule_with_warmup

from src.dataloader import get_dataloader
from src.model import SlotFillingModel
from config import DataTrainingArguments, ModelArguments
from trainer import train, eval
from src.utils import log_params

import sys
import os
import time
import json

def main():
    # parse arguments 
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # For logging
    current_time = time.localtime()
    current_time = f"{current_time.tm_year}_{current_time.tm_mon}_{current_time.tm_mday}_{current_time.tm_hour}_{current_time.tm_min}_{current_time.tm_sec}"
    
    save_path = f'{training_args.output_dir}/{data_args.target_domain}/Sample{data_args.n_samples}/'
    # log_path = f'{training_args.output_dir}/{data_args.target_domain}/Sample{data_args.n_samples}/'
    # model_path = f'{training_args.output_dir}/model/{data_args.target_domain}/Sample{data_args.n_samples}/'
    log_dict = {}
    
    log_params(log_dict, [model_args, data_args, training_args])

    # load pretrained BERT and define model 
    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', use_fast=True)
    model = nn.DataParallel(SlotFillingModel(model_args).cuda()) if torch.cuda.is_available() else SlotFillingModel(model_args)
    
    # get dataloader
    dataloader_train, dataloader_val, dataloader_test = get_dataloader(
                                            data_args.target_domain, 
                                            training_args.per_device_train_batch_size, 
                                            data_args.n_samples, 
                                            data_args.dataset_path,
                                            tokenizer,)

    if data_args.run_mode == 'train':
        print("Training mode...")
        # loss function, optimizer, ...
        optim = AdamW(model.parameters(), lr=training_args.learning_rate, correct_bias=True)

        scheduler = get_linear_schedule_with_warmup(optim, num_warmup_steps=training_args.warmup_steps, num_training_steps=training_args.max_steps)

        os.makedirs(save_path, exist_ok=True)
        
        print(f'Target Domain: {data_args.target_domain}\tN Samples: {data_args.n_samples}')

        best_step, best_f1 = train(model=model, 
                                    dataloader_train=dataloader_train, 
                                    dataloader_val=dataloader_val, 
                                    optim=optim, 
                                    scheduler=scheduler, 
                                    eval_steps=training_args.eval_steps,
                                    total_steps=training_args.max_steps,
                                    early_stopping_patience=data_args.early_stopping_patience,
                                    model_save_path=save_path,
                                    log_dict=log_dict)

        print("Training finished.")
        print(f"Best validation f1 score {best_f1: .2f} at training step {best_step}")

        with open(save_path + 'log.json', 'w') as json_out:
            json.dump(log_dict, json_out, indent=4)

    
    elif data_args.run_mode == 'test':
        print("Test mode...")
        # Prediction / Test
        model.load_state_dict(torch.load(save_path+f"best-model-parameters.pt"))
        results = eval(model, dataloader_test, data_args.target_domain, tokenizer, save_path+"test_output.json")
        print(f"F1 Score at prediction: {results['fb1']}")

        log_dict['test_result'] = results['fb1']
        

    else:
        print("Invalid input: option \"run_mode\" got wrong value.")
    
    return 


if __name__=="__main__":
    main()