train.py

import logging
import os
import sys
import pandas as pd
import torch

from typing import List, Callable, NoReturn, NewType, Any
import dataclasses
from datasets import (
    load_metric,
    load_from_disk,
    Dataset,
    DatasetDict,
    Features,
    Sequence,
    Value,
    concatenate_datasets,
)

from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    EarlyStoppingCallback,
)

from transformers import (
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    TrainingArguments,
    set_seed,
)

from tokenizers import Tokenizer
from tokenizers.models import WordPiece

from utils_qa import postprocess_qa_predictions, check_no_error
from trainer_qa import QuestionAnsweringTrainer

from arguments import (
    ModelArguments,
    DataTrainingArguments,
)

from elastic_retrieval import SparseRetrieval
import wandb
from elastic_setting import preprocess

logger = logging.getLogger(__name__)


def main():
    # 가능한 arguments 들은 ./arguments.py 나 transformer package 안의 src/transformers/training_args.py 에서 확인 가능합니다.
    # --help flag 를 실행시켜서 확인할 수 도 있습니다.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments))
    model_args, data_args = parser.parse_args_into_dataclasses()
    print(model_args.model_name_or_path)

    # [참고] argument를 manual하게 수정하고 싶은 경우에 아래와 같은 방식을 사용할 수 있습니다
    # training_args.per_device_train_batch_size = 4
    # print(training_args.per_device_train_batch_size)
    # 5-fold 데이터를 불러옵니다.
    fold1 = pd.read_csv("./fold1.csv")
    fold2 = pd.read_csv("./fold2.csv")
    fold3 = pd.read_csv("./fold3.csv")
    fold4 = pd.read_csv("./fold4.csv")
    fold5 = pd.read_csv("./fold5.csv")

    fold1["answers"] = fold1["answers"].apply(eval)
    fold2["answers"] = fold2["answers"].apply(eval)
    fold3["answers"] = fold3["answers"].apply(eval)
    fold4["answers"] = fold4["answers"].apply(eval)
    fold5["answers"] = fold5["answers"].apply(eval)

    folds = [fold1, fold2, fold3, fold4, fold5]

    # range안의 숫자를 조절하여 특정 fold만 학습할 수 있습니다. fold로 들어갈 수 있는 숫자는 1~5입니다.(range(1,6))
    for fold in range(1, 6):
        training_args = TrainingArguments(
            do_train=True,
            output_dir="./models/train_dataset_ng5_fold" + str(fold),
            overwrite_output_dir=True,
            evaluation_strategy="steps",
            per_device_train_batch_size=16,
            per_device_eval_batch_size=16,
            gradient_accumulation_steps=2,
            learning_rate=1e-5,
            num_train_epochs=8,
            warmup_ratio=0.1,
            logging_strategy="steps",
            logging_steps=100,
            save_strategy="steps",
            save_steps=300,
            save_total_limit=1,
            seed=42,
            eval_steps=300,
            metric_for_best_model="exact_match",
            load_best_model_at_end=True,
        )
        print(f"model is from {model_args.model_name_or_path}")
        print(f"data is from {data_args.dataset_name}")

        # 모델을 초기화하기 전에 난수를 고정합니다.
        set_seed(training_args.seed)

        # logging 설정
        logging.basicConfig(
            format="%(asctime)s - %(levelname)s - %(name)s -    %(message)s",
            datefmt="%m/%d/%Y %H:%M:%S",
            handlers=[logging.StreamHandler(sys.stdout)],
        )

        # verbosity 설정 : Transformers logger의 정보로 사용합니다 (on main process only)
        logger.info("Training/evaluation parameters %s", training_args)

        df = pd.concat([folds[i] for i in range(5) if i + 1 != fold], ignore_index=True)

        f = Features(
            {
                "answers": Sequence(
                    feature={
                        "text": Value(dtype="string", id=None),
                        "answer_start": Value(dtype="int32", id=None),
                    },
                    length=-1,
                    id=None,
                ),
                "context": Value(dtype="string", id=None),
                "id": Value(dtype="string", id=None),
                "question": Value(dtype="string", id=None),
            }
        )
        datasets = DatasetDict(
            {
                "train": Dataset.from_pandas(df, features=f),
                "validation": Dataset.from_pandas(folds[fold - 1], features=f),
            }
        )

        """
        # datasets = load_from_disk(data_args.dataset_name)
        datasets.save_to_disk("./fold_dataset/")
        datasets = load_from_disk("./fold_dataset/")
        """
        print(datasets)

        # AutoConfig를 이용하여 pretrained model 과 tokenizer를 불러옵니다.
        # argument로 원하는 모델 이름을 설정하면 옵션을 바꿀 수 있습니다.
        config = AutoConfig.from_pretrained(
            model_args.config_name
            if model_args.config_name is not None
            else model_args.model_name_or_path,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.tokenizer_name
            if model_args.tokenizer_name is not None
            else model_args.model_name_or_path,
            # 'use_fast' argument를 True로 설정할 경우 rust로 구현된 tokenizer를 사용할 수 있습니다.
            # False로 설정할 경우 python으로 구현된 tokenizer를 사용할 수 있으며,
            # rust version이 비교적 속도가 빠릅니다.
            use_fast=True,
        )
        model = AutoModelForQuestionAnswering.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
        )
        print(
            type(training_args),
            type(model_args),
            type(datasets),
            type(tokenizer),
            type(model),
        )

        # do_train mrc model 혹은 do_eval mrc model
        if training_args.do_train or training_args.do_eval:
            run = wandb.init(
                project="mrc",
                entity="quarter100",
                name="fold" + str(fold),
                group="ng5only_fold",
            )
            run_mrc(data_args, training_args, model_args, datasets, tokenizer, model)
            run.finish()


def run_mrc(
    data_args: DataTrainingArguments,
    training_args: TrainingArguments,
    model_args: ModelArguments,
    datasets: DatasetDict,
    tokenizer,
    model,
) -> NoReturn:

    # dataset을 전처리합니다.
    # training과 evaluation에서 사용되는 전처리는 아주 조금 다른 형태를 가집니다.
    if training_args.do_train:
        column_names = datasets["train"].column_names
    else:
        column_names = datasets["validation"].column_names

    question_column_name = "question" if "question" in column_names else column_names[0]
    context_column_name = "context" if "context" in column_names else column_names[1]
    answer_column_name = "answers" if "answers" in column_names else column_names[2]

    # Padding에 대한 옵션을 설정합니다.
    # (question|context) 혹은 (context|question)로 세팅 가능합니다.
    pad_on_right = tokenizer.padding_side == "right"

    # 오류가 있는지 확인합니다.
    last_checkpoint, max_seq_length = check_no_error(
        data_args, training_args, datasets, tokenizer
    )
    """
    qg_df = pd.read_pickle("../data/delete_qg_sort.pkl")
    qg_df = qg_df.iloc[1:1200]
    qg_dataset = Dataset.from_pandas(qg_df)
    qg_dataset.save_to_disk("../data/qg_dataset/")
    """
    ES_retrieval = SparseRetrieval()

    # Train preprocessing / 전처리를 진행합니다.
    def prepare_train_features(examples):
        # truncation과 padding(length가 짧을때만)을 통해 toknization을 진행하며, stride를 이용하여 overflow를 유지합니다.
        # 각 example들은 이전의 context와 조금씩 겹치게됩니다.
        # 원본 데이터 context 전처리 및 그에 따른 answer_position 이동
        for i in range(len(examples[context_column_name])):
            context = examples[context_column_name][i]
            answer = examples[answer_column_name][i]
            answer_start = answer["answer_start"][0]
            answer_text = answer["text"][0]
            answer_end = answer_start + len(answer_text)
            context_pre = context[:answer_start]
            context_post = context[answer_end:]
            context_pre = preprocess(context_pre)
            context_post = preprocess(context_post)
            new_answer_start = len(context_pre)
            examples[context_column_name][i] = context_pre + answer_text + context_post
            examples[answer_column_name][i] = {
                "answer_start": [new_answer_start],
                "text": [answer_text],
            }

        tokenized_examples = tokenizer(
            examples[question_column_name if pad_on_right else context_column_name],
            examples[context_column_name if pad_on_right else question_column_name],
            truncation="only_second" if pad_on_right else "only_first",
            max_length=max_seq_length,
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            return_token_type_ids=False,  # roberta모델을 사용할 경우 False, bert를 사용할 경우 True로 표기해야합니다.
            padding="max_length" if data_args.pad_to_max_length else False,
        )

        # 길이가 긴 context가 등장할 경우 truncate를 진행해야하므로, 해당 데이터셋을 찾을 수 있도록 mapping 가능한 값이 필요합니다.
        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
        # token의 캐릭터 단위 position를 찾을 수 있도록 offset mapping을 사용합니다.
        # start_positions과 end_positions을 찾는데 도움을 줄 수 있습니다.
        offset_mapping = tokenized_examples.pop("offset_mapping")

        # 데이터셋에 "start position", "enc position" label을 부여합니다.
        tokenized_examples["start_positions"] = []
        tokenized_examples["end_positions"] = []

        for i, offsets in enumerate(offset_mapping):
            input_ids = tokenized_examples["input_ids"][i]
            cls_index = input_ids.index(tokenizer.cls_token_id)  # cls index

            # sequence id를 설정합니다 (to know what is the context and what is the question).
            sequence_ids = tokenized_examples.sequence_ids(i)

            # 하나의 example이 여러개의 span을 가질 수 있습니다.
            sample_index = sample_mapping[i]
            answers = examples[answer_column_name][sample_index]

            # answer가 없을 경우 cls_index를 answer로 설정합니다(== example에서 정답이 없는 경우 존재할 수 있음).
            if len(answers["answer_start"]) == 0:
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                # text에서 정답의 Start/end character index
                start_char = answers["answer_start"][0]
                end_char = start_char + len(answers["text"][0])

                # text에서 current span의 Start token index
                token_start_index = 0
                while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                    token_start_index += 1

                # text에서 current span의 End token index
                token_end_index = len(input_ids) - 1
                while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                    token_end_index -= 1

                # 정답이 span을 벗어났는지 확인합니다(정답이 없는 경우 CLS index로 label되어있음).
                if not (
                    offsets[token_start_index][0] <= start_char
                    and offsets[token_end_index][1] >= end_char
                ):
                    tokenized_examples["start_positions"].append(cls_index)
                    tokenized_examples["end_positions"].append(cls_index)
                else:
                    # token_start_index 및 token_end_index를 answer의 끝으로 이동합니다.
                    # Note: answer가 마지막 단어인 경우 last offset을 따라갈 수 있습니다(edge case).
                    while (
                        token_start_index < len(offsets)
                        and offsets[token_start_index][0] <= start_char
                    ):
                        token_start_index += 1
                    tokenized_examples["start_positions"].append(token_start_index - 1)
                    while offsets[token_end_index][1] >= end_char:
                        token_end_index -= 1
                    tokenized_examples["end_positions"].append(token_end_index + 1)
        return tokenized_examples

    # negative sampling
    def prepare_train_features_ng(examples):
        x = Dataset.from_dict(examples)
        negative_df = ES_retrieval.retrieve_ES(
            x, topk=data_args.ng_top_k_retrieval, ner_path="./train_tagged.csv"
        )
        negative_query = negative_df["question"]
        negative_contexts = negative_df["context"]
        negative_gt_contexts = negative_df["original_context"]
        nq_final = []
        nc_final = []
        for i in range(len(negative_query)):
            temp_q = [
                negative_query[i] for _ in range((data_args.ng_top_k_retrieval - 1))
            ]
            nq_final.extend(temp_q)
            if negative_gt_contexts[i] in negative_contexts[i]:
                negative_contexts[i].remove(negative_gt_contexts[i])
                temp_c = negative_contexts[i]
                nc_final.extend(temp_c)
            else:
                temp_c = negative_contexts[i][: (data_args.ng_top_k_retrieval - 1)]
                nc_final.extend(temp_c)
        assert len(nq_final) == len(
            nc_final
        ), f"nq_final length {len(nq_final)} should be same as nc_final {len(nc_final)}"

        tokenized_examples_ng = tokenizer(
            nq_final if pad_on_right else nc_final,
            nc_final if pad_on_right else nq_final,
            truncation="only_second" if pad_on_right else "only_first",
            max_length=max_seq_length,
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            return_token_type_ids=False,  # roberta모델을 사용할 경우 False, bert를 사용할 경우 True로 표기해야합니다.
            padding="max_length" if data_args.pad_to_max_length else False,
        )
        tokenized_examples_ng.pop("overflow_to_sample_mapping")
        tokenized_examples_ng.pop("offset_mapping")
        tokenized_examples_ng["start_positions"] = []
        tokenized_examples_ng["end_positions"] = []

        for i in range(len(tokenized_examples_ng["input_ids"])):
            tokenized_examples_ng["start_positions"].append(0)
            tokenized_examples_ng["end_positions"].append(0)
        return tokenized_examples_ng

    if training_args.do_train:
        if "train" not in datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = datasets["train"]

        # dataset에서 train feature를 생성합니다.
        train_dataset_ps = train_dataset.map(
            prepare_train_features,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=data_args.overwrite_cache,
        )
        train_dataset_ng = train_dataset.map(
            prepare_train_features_ng,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=data_args.overwrite_cache,
        )
        """
        qg_dataset = load_from_disk("../data/qg_dataset/")
        train_dataset_qg = qg_dataset.map(
            prepare_train_features,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=qg_dataset.column_names,
            load_from_cache_file=data_args.overwrite_cache,
        )
        """
        train_dataset = concatenate_datasets(
            [
                train_dataset_ps.flatten_indices(),
                train_dataset_ng.flatten_indices(),
                #train_dataset_qg.flatten_indices(),
            ]
        )
    print("train_dataset length : ", len(train_dataset))

    # Validation preprocessing
    def prepare_validation_features(examples):
        # truncation과 padding(length가 짧을때만)을 통해 toknization을 진행하며, stride를 이용하여 overflow를 유지합니다.
        # 각 example들은 이전의 context와 조금씩 겹치게됩니다.
        tokenized_examples = tokenizer(
            examples[question_column_name if pad_on_right else context_column_name],
            examples[context_column_name if pad_on_right else question_column_name],
            truncation="only_second" if pad_on_right else "only_first",
            max_length=max_seq_length,
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            return_token_type_ids=False,  # roberta모델을 사용할 경우 False, bert를 사용할 경우 True로 표기해야합니다.
            padding="max_length" if data_args.pad_to_max_length else False,
        )

        # 길이가 긴 context가 등장할 경우 truncate를 진행해야하므로, 해당 데이터셋을 찾을 수 있도록 mapping 가능한 값이 필요합니다.
        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")

        # evaluation을 위해, prediction을 context의 substring으로 변환해야합니다.
        # corresponding example_id를 유지하고 offset mappings을 저장해야합니다.
        tokenized_examples["example_id"] = []

        for i in range(len(tokenized_examples["input_ids"])):
            # sequence id를 설정합니다 (to know what is the context and what is the question).
            sequence_ids = tokenized_examples.sequence_ids(i)
            context_index = 1 if pad_on_right else 0

            # 하나의 example이 여러개의 span을 가질 수 있습니다.
            sample_index = sample_mapping[i]
            tokenized_examples["example_id"].append(examples["id"][sample_index])

            # Set to None the offset_mapping을 None으로 설정해서 token position이 context의 일부인지 쉽게 판별 할 수 있습니다.
            tokenized_examples["offset_mapping"][i] = [
                (o if sequence_ids[k] == context_index else None)
                for k, o in enumerate(tokenized_examples["offset_mapping"][i])
            ]
        return tokenized_examples

    if training_args.do_eval:
        eval_dataset = datasets["validation"]
        for i in range(len(eval_dataset[context_column_name])):
            context = eval_dataset[context_column_name][i]
            answer = eval_dataset[answer_column_name][i]
            answer_start = answer["answer_start"][0]
            answer_text = answer["text"][0]
            answer_end = answer_start + len(answer_text)
            context_pre = context[:answer_start]
            context_post = context[answer_end:]
            context_pre = preprocess(context_pre)
            context_post = preprocess(context_post)
            new_answer_start = len(context_pre)
            eval_dataset[context_column_name][i] = (
                context_pre + answer_text + context_post
            )
            eval_dataset[answer_column_name][i] = {
                "answer_start": [new_answer_start],
                "text": [answer_text],
            }

        # Validation Feature 생성
        eval_dataset_tokenized = eval_dataset.map(
            prepare_validation_features,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=data_args.overwrite_cache,
        )
    print("valid_dataset length : ", len(eval_dataset_tokenized))

    # Data collator
    # flag가 True이면 이미 max length로 padding된 상태입니다.
    # 그렇지 않다면 data collator에서 padding을 진행해야합니다.
    data_collator = DataCollatorWithPadding(
        tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None
    )

    # Post-processing:
    def post_processing_function(examples, features, predictions, training_args):
        # Post-processing: start logits과 end logits을 original context의 정답과 match시킵니다.
        predictions = postprocess_qa_predictions(
            examples=examples,
            features=features,
            predictions=predictions,
            max_answer_length=data_args.max_answer_length,
            output_dir=training_args.output_dir,
        )
        # Metric을 구할 수 있도록 Format을 맞춰줍니다.
        formatted_predictions = [
            {"id": k, "prediction_text": v} for k, v in predictions.items()
        ]
        if training_args.do_predict:
            return formatted_predictions

        elif training_args.do_eval:
            references = [
                {"id": ex["id"], "answers": ex[answer_column_name]}
                for ex in datasets["validation"]
            ]
            return EvalPrediction(
                predictions=formatted_predictions, label_ids=references
            )

    metric = load_metric("squad")

    def compute_metrics(p: EvalPrediction):
        result = metric.compute(predictions=p.predictions, references=p.label_ids)
        result["eval_exact_match"] = result["exact_match"]
        del result["exact_match"]
        result["eval_f1"] = result["f1"]
        del result["f1"]
        return result

    # Trainer 초기화
    trainer = QuestionAnsweringTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset_tokenized if training_args.do_eval else None,
        eval_examples=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        post_process_function=post_processing_function,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
    )

    # Training
    if training_args.do_train:
        if last_checkpoint is not None:
            checkpoint = last_checkpoint
        elif os.path.isdir(model_args.model_name_or_path):
            checkpoint = model_args.model_name_or_path
        else:
            checkpoint = None
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        metrics["train_samples"] = len(train_dataset)

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")

        with open(output_train_file, "w") as writer:
            logger.info("***** Train results *****")
            for key, value in sorted(train_result.metrics.items()):
                logger.info(f"  {key} = {value}")
                writer.write(f"{key} = {value}\n")

        # State 저장
        trainer.state.save_to_json(
            os.path.join(training_args.output_dir, "trainer_state.json")
        )

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate()

        metrics["eval_samples"] = len(eval_dataset_tokenized)

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)


if __name__ == "__main__":
    main()