train.py

import pickle as pickle
import os
import pandas as pd
import torch
import sklearn
import numpy as np
from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score
from transformers import (
  AutoTokenizer, 
  AutoConfig, 
  AutoModelForSequenceClassification, 
  Trainer, 
  TrainingArguments, 
  RobertaConfig,
  RobertaTokenizer, 
  RobertaForSequenceClassification, 
  BertTokenizer,
  EarlyStoppingCallback
)    
from load_data import *
import wandb
import json
import random
from test_recording import *
from datasets import load_dataset, load_metric
from sklearn.metrics import classification_report
from training_loss import CustomTrainer

def seed_everything(seed: int = 42):
    """Random seed(Reproducibility)"""
    random.seed(seed)                              
    np.random.seed(seed)                           
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.manual_seed(seed)                        
    torch.cuda.manual_seed(seed)  # type: ignore    
    torch.backends.cudnn.deterministic = True  # type: ignore
    torch.backends.cudnn.benchmark = False  # type: ignore

def ruw_re_micro_f1(preds, labels):
    """Russia-Ukraine War RE micro f1 (except no_relation)"""
    label_list = ['no_relation', 'per:title', 'org:member_of', 
        'org:alternate_names', 'org:top_members/employees', 'org:hostile', 
        'org:friendly', 'org:property', 'eve:place', 'eve:date'] 

    no_relation_label_idx = label_list.index("no_relation")
    label_indices = list(range(len(label_list)))
    label_indices.remove(no_relation_label_idx)
    return sklearn.metrics.f1_score(labels, preds, average="micro", labels=label_indices) * 100.0

def ruw_re_auprc(probs, labels):
    """Russia-Ukraine War RE AUPRC (with no_relation)"""
    labels = np.eye(10)[labels]

    score = np.zeros((10,))
    for c in range(10):
        targets_c = labels.take([c], axis=1).ravel()
        preds_c = probs.take([c], axis=1).ravel()
        precision, recall, _ = sklearn.metrics.precision_recall_curve(targets_c, preds_c)
        score[c] = sklearn.metrics.auc(recall, precision)
    return np.average(score) * 100.0

def compute_metrics(pred):
  """ validation을 위한 metrics function """
  labels = pred.label_ids
  preds = pred.predictions.argmax(-1)
  probs = pred.predictions
  label_list = ['no_relation', 'per:title', 'org:member_of', 
        'org:alternate_names', 'org:top_members/employees', 'org:hostile', 
        'org:friendly', 'org:property', 'eve:place', 'eve:date']
  print(classification_report(labels, preds, target_names=label_list))

  # calculate accuracy using sklearn's function
  f1 = ruw_re_micro_f1(preds, labels)
  auprc = ruw_re_auprc(probs, labels)
  acc = accuracy_score(labels, preds) # 리더보드 평가에는 포함되지 않습니다.

  return {
      'micro f1 score': f1,
      'auprc' : auprc,
      'accuracy': acc,
  }

def label_to_num(label):
  num_label = []
  with open('ruw_label_to_num.pkl', 'rb') as f:
    dict_label_to_num = pickle.load(f)
  for v in label:
    num_label.append(dict_label_to_num[v])
  return num_label

def sort_by_len(dataset):
  """ dataframe을 sentence 길이로 정렬해 반환 """

  sent_len = list(dataset['sentence'])
  each_sent_len = [len(sent) for sent in sent_len]

  dataset['length'] = each_sent_len
  dataset = dataset.sort_values(by = ['length'])
  return dataset

def train():
  # load_parameter: tokenizer, sentence preprocessing
  with open("config.json","r") as js:
    config = json.load(js)
    load_model = config['model_name']        # model
    filter = config['sentence_filter']       # sentence_filter
    marking_mode = config['marking_mode']    # marking_mode
    tokenize_mode = config['tokenize_mode'] # tokenize_function
    wandb_name = config['test_name']
    loss_name = config['loss_name']  #loss_name
    train_dataloader = config['train_dataloader']
  
  # load model and tokenizer  # MODEL_NAME = "bert-base-uncased"
  MODEL_NAME = load_model
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
  print("#################################################################################################################### \n",
        f"Model_name: {MODEL_NAME}, Filter: {filter}, Marking_mode: {marking_mode}, Tokenized_function: {tokenize_mode}\n",
        "#################################################################################################################### \n")

  # load dataset
  dataset_dir = "ruw_data/train.csv"
  train_dataset, dev_dataset = load_data(dataset_dir, train=True, filter=filter, marking_mode=marking_mode)
  # train_dataset, dev_dataset = load_aug_data(dataset_dir, train=True, filter=filter, marking_mode=marking_mode, aug_type="swap", save=True)  # augmentation 사용시
  if train_dataloader == "sequential":
   train_dataset = sort_by_len(train_dataset)  # 문장 길이로 정렬
  train_label = label_to_num(train_dataset['label'].values)
  dev_label = label_to_num(dev_dataset['label'].values)
  
  # add vocab (special tokens)
  with open("marking_mode_tokens.json","r") as json_file:
    mode2special_token = json.load(json_file)
  add_token_num = 0
  if marking_mode != "normal" and  marking_mode != "typed_entity_punc":
    add_token_num += tokenizer.add_special_tokens({"additional_special_tokens":mode2special_token[marking_mode]})
  
  # tokenizing dataset
  tokenized_train = tokenized_dataset(train_dataset, tokenizer, tokenize_mode)
  tokenized_dev = tokenized_dataset(dev_dataset, tokenizer, tokenize_mode)

  # make dataset for pytorch.
  RE_train_dataset = RE_Dataset(tokenized_train, train_label)
  RE_dev_dataset = RE_Dataset(tokenized_dev, dev_label)

  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

  print(device)
  # setting model hyperparameter
  model_config =  AutoConfig.from_pretrained(MODEL_NAME)
  model_config.num_labels = 10  # klue - 30

  model =  AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, config=model_config)
  #resize models vocab_size(add add_token_num) 
  model.resize_token_embeddings(tokenizer.vocab_size + add_token_num)
  # print(model.config)
  model.parameters
  model.to(device)
  
  project = "Russia-Ukraine-War"  # W&B Projects
  entity_name = "hannayeoniee"
  display_name = wandb_name  # Model_name displayed in W&B Projects
  wandb.init(project=project, entity=entity_name, name=display_name)
  
  # 사용한 option 외에도 다양한 option들이 있습니다.
  # https://huggingface.co/transformers/main_classes/trainer.html#trainingarguments 참고해주세요.
  # 현재 TrainingArguments에 정의된 디폴트 파라미터와 config.json 파일에서의 옵션값은 SOTA 모델에 사용된 값으로 성능을 재현할 수 있습니다 
  training_args = TrainingArguments(
    output_dir='./results',          # output directory
    save_total_limit=10,              # number of total save model.
    save_steps=50,                 # model saving step.
    num_train_epochs=10,              # total number of training epochs
    learning_rate=3e-5,               # learning_rate
    per_device_train_batch_size=32,  # batch size per device during training
    per_device_eval_batch_size=16,   # batch size for evaluation
    # added max_length in load_data.py
    warmup_ratio = 0.1,  # defalut 0
    adam_epsilon = 1e-6, # default 1e-8
    # warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
    logging_steps=100,              # log saving step.
    evaluation_strategy='steps', # evaluation strategy to adopt during training
                                # `no`: No evaluation during training.
                                # `steps`: Evaluate every `eval_steps`.
                                # `epoch`: Evaluate every end of epoch.
    eval_steps = 50,            # evaluation step. (default:500)
    load_best_model_at_end = True,
    metric_for_best_model = 'micro f1 score',
    report_to="wandb",  # enable logging to W&B
    fp16 = True,        # whether to use 16bit (mixed) precision training
    fp16_opt_level = 'O1' # choose AMP optimization level (AMP Option:'O1' , 'O2')(FP32: 'O0')
  )
  # save test result 
  save_record(config, training_args)

  # use custom trainer for using custom training loss
  trainer = CustomTrainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=RE_train_dataset,         # training dataset
    eval_dataset=RE_dev_dataset,             # evaluation dataset
    compute_metrics=compute_metrics,         # define metrics function
    # callbacks=[EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.0)], #EarlyStopping callbacks
    original_dataset = train_dataset,
    device = device,
    loss_name = loss_name,                 # set loss for backpropagation
    train_dataloader = train_dataloader
  )

  # train model
  trainer.train()
  model.save_pretrained('./best_model')
def main():
  train()

if __name__ == '__main__':
  seed_everything(42)
  main()