inference.py

##################
# Import modules #
##################

from typing import Callable, List, Dict, NoReturn, Tuple, Optional
import numpy as np
import os
import argparse

from datasets import (
    load_metric,
    load_from_disk,
    Sequence,
    Value,
    Features,
    Dataset,
    DatasetDict,
)

from transformers import AutoConfig, AutoModelForQuestionAnswering, AutoTokenizer, RobertaModel, AutoConfig

from transformers import (
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    TrainingArguments,
    set_seed,
)
from model.Reader.RobertaCnn import RobertaCNNForQuestionAnswering
from utils.utils_qa import postprocess_qa_predictions, check_no_error
from model.Reader.trainer_qa import QuestionAnsweringTrainer
from model.Retrieval.retrieval import DenseRetrieval, SparseRetrieval, JointRetrieval

from utils.arguments import (
    ModelArguments,
    DataArguments,
    DenseTrainingArguments,

    inference_config_setting,
    wandb_config_setting,
    
    INFERENCE_DIR,
    CONFIG_DIR,
    LOG_DIR,
)
from utils.logger import get_logger


########################
# Set global variables #
########################

logger = None
CUSTOM_MODEL_NAMES = {
    "RobertaCnn":RobertaCNNForQuestionAnswering,
}

#######################
# Classes & Functions #
#######################

def main():
    
    global logger

    # Load config json
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config_file_path", help="Configure Json path")
    parser.add_argument("-l", "--log_file_path", default="reader_train.log", help="Logger file path")
    parser.add_argument("-n", "--inference_name", default=None, help="Inference file directory name")
    parser.add_argument("-m", "--model_name_or_path", default=None, help="Reader model path for inference")
    parser.add_argument("--do_predict",action="store_true")
    
    config = parser.parse_args()

    assert config.inference_name, "Output 파일 이름을 설정해 주세요"

    config.config_file_path = os.path.join(CONFIG_DIR, config.config_file_path) 
    config.log_file_path = os.path.join(LOG_DIR, config.log_file_path)
    config.inference_name = os.path.join(INFERENCE_DIR, config.inference_name)

    model_args, data_args, dense_args, training_args =\
        inference_config_setting(config)

    training_args.do_train = True
    training_args.do_predict = config.do_predict

    logger = get_logger("logs/inference.log")

    logger.info(f"model is from {model_args.model_name_or_path}")
    logger.info(f"data is from {data_args.dataset_name}")
    logger.info("Training/evaluation parameters %s", training_args)

    # Set random seed
    set_seed(training_args.seed)

    datasets = load_from_disk(data_args.dataset_name)
    logger.info(datasets)

    # Load Config & tokenizer
    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=True,
    )

    # Load Model
    if model_args.model_name_or_path in CUSTOM_MODEL_NAMES:
        model = CUSTOM_MODEL_NAMES[model_args.model_name_or_path].from_pretrained(
            model_args.model_name_or_path,
            config=config,
        )
    else:
        model = AutoModelForQuestionAnswering.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
        )

    # True일 경우 : run passage retrieval
    if data_args.eval_retrieval:
        if data_args.kind_of_retrieval == "Sparse":
            datasets = run_sparse_retrieval(
                tokenizer.tokenize,
                datasets,
                training_args,
                data_args,
            )
        elif data_args.kind_of_retrieval == "Dense":
            datasets = run_dense_retrieval(
                datasets,
                training_args,
                data_args,
                dense_args,
            )
        elif data_args.kind_of_retrieval == "Joint":
            datasets = run_joint_retrieval(
                tokenizer.tokenize,
                datasets,
                training_args,
                data_args,
                dense_args,
            )

    # eval or predict mrc model
    if training_args.do_eval or training_args.do_predict:
        run_mrc(data_args, training_args, model_args, datasets, tokenizer, model)

def run_joint_retrieval(
    tokenize_fn: Callable[[str], List[str]],
    datasets: DatasetDict,
    training_args: TrainingArguments,
    dense_args: DenseTrainingArguments,
    data_args: DataArguments,
    data_path: str = "../data",
    context_path: str = "wikipedia_documents.json",
    embedding_form : Optional[str] = "BM25"
) -> DatasetDict:
    
    p_tokenizer = AutoTokenizer.from_pretrained('klue/roberta-small')
    q_tokenizer = AutoTokenizer.from_pretrained('klue/roberta-small')
    
    p_encoder = RobertaModel.from_pretrained(dense_args.dense_passage_retrieval_name).to('cuda')
    q_encoder = RobertaModel.from_pretrained(dense_args.dense_question_retrieval_name).to('cuda')
    
    retriever = JointRetrieval(
        sparse_tokenize_fn = tokenize_fn,
        dense_tokenizer = (p_tokenizer, q_tokenizer),
        encoders = (p_encoder, q_encoder),
        data_path = data_path,
        context_path = context_path,
        embedding_form = embedding_form
    )

    df = retriever.retrieve(datasets["validation"], topk=data_args.top_k_retrieval)

    # test data 에 대해선 정답이 없으므로 id question context 로만 데이터셋이 구성됩니다.
    if training_args.do_predict:
        f = Features(
            {
                "context": Value(dtype="string", id=None),
                "id": Value(dtype="string", id=None),
                "question": Value(dtype="string", id=None),
            }
        )

    # train data 에 대해선 정답이 존재하므로 id question context answer 로 데이터셋이 구성됩니다.
    elif training_args.do_eval:
        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({"validation": Dataset.from_pandas(df, features=f)})
    return datasets


def run_sparse_retrieval(
    tokenize_fn: Callable[[str], List[str]],
    datasets: DatasetDict,
    training_args: TrainingArguments,
    data_args: DataArguments,
    data_path: str = "../data",
    context_path: str = "wikipedia_documents.json",
) -> DatasetDict:

    # Query에 맞는 Passage들을 Retrieval 합니다.
    retriever = SparseRetrieval(
        tokenize_fn=tokenize_fn,
        data_path=data_path,
        context_path=context_path,
        embedding_form="ES",
    )
    retriever.get_sparse_embedding()

    if data_args.use_faiss:
        retriever.build_faiss(num_clusters=data_args.num_clusters)
        df = retriever.retrieve_faiss(
            datasets["validation"], topk=data_args.top_k_retrieval
        )
    else:
        df = retriever.retrieve(datasets["validation"], topk=data_args.top_k_retrieval)

    # test data 에 대해선 정답이 없으므로 id question context 로만 데이터셋이 구성됩니다.
    if training_args.do_predict:
        f = Features(
            {
                "context": Value(dtype="string", id=None),
                "id": Value(dtype="string", id=None),
                "question": Value(dtype="string", id=None),
            }
        )

    # train data 에 대해선 정답이 존재하므로 id question context answer 로 데이터셋이 구성됩니다.
    elif training_args.do_eval:
        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({"validation": Dataset.from_pandas(df, features=f)})
    return datasets


def run_dense_retrieval(
    datasets: DatasetDict,
    training_args: TrainingArguments,
    data_args: DataArguments,
    dense_args: DenseTrainingArguments,
    data_path: str = "./data",
    context_path: str = "wikipedia_documents.json",
) -> DatasetDict:
    ## 1. p 인코더, q 인코더 불러오기
        # Query에 맞는 Passage들을 Retrieval 합니다.
    p_tokenizer = AutoTokenizer.from_pretrained('Huffon/sentence-klue-roberta-base')#'klue/roberta-small')
    q_tokenizer = AutoTokenizer.from_pretrained('Huffon/sentence-klue-roberta-base')#'klue/roberta-small')
    
    p_encoder = RobertaModel.from_pretrained(dense_args.dense_passage_retrieval_name).to('cuda')
    q_encoder = RobertaModel.from_pretrained(dense_args.dense_question_retrieval_name).to('cuda')
    retriever = DenseRetrieval(
        tokenizers=(p_tokenizer, q_tokenizer), encoders= (p_encoder, q_encoder), data_path=data_path, context_path=context_path)

    ## 2. passage embeddings 구하기
    retriever.get_dense_passage_embedding()
    
    del p_encoder # 메모리 확보
    
    ## 3. 각 쿼리 임베딩에 따른 passage 구하기
    df = retriever.retrieve(q_encoder, datasets["validation"], topk=data_args.top_k_retrieval)

    ## 4. 반환하기
    # test data 에 대해선 정답이 없으므로 id question context 로만 데이터셋이 구성됩니다.
    if training_args.do_predict:
        f = Features(
            {
                "context": Value(dtype="string", id=None),
                "id": Value(dtype="string", id=None),
                "question": Value(dtype="string", id=None),
            }
        )

    # train data 에 대해선 정답이 존재하므로 id question context answer 로 데이터셋이 구성됩니다.
    elif training_args.do_eval:
        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({"validation": Dataset.from_pandas(df, features=f)})
    return datasets

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

    # eval 혹은 prediction에서만 사용함
    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
    )

    # 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])

            # context의 일부가 아닌 offset_mapping을 None으로 설정하여 토큰 위치가 컨텍스트의 일부인지 여부를 쉽게 판별할 수 있습니다.
            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

    eval_dataset = datasets["validation"]

    # Validation Feature 생성
    eval_dataset = eval_dataset.map(
        prepare_validation_features,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    # 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: Tuple[np.ndarray, np.ndarray],
        training_args: TrainingArguments,
    ) -> EvalPrediction:
        # 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) -> Dict:
        return metric.compute(predictions=p.predictions, references=p.label_ids)

    logger.info("init trainer...")
    # logger.info(model)
    # exit(0)
    # Trainer 초기화
    trainer = QuestionAnsweringTrainer(
        model=model,
        args=training_args,
        train_dataset=None,
        eval_dataset=eval_dataset,
        eval_examples=datasets["validation"],
        tokenizer=tokenizer,
        data_collator=data_collator,
        post_process_function=post_processing_function,
        compute_metrics=compute_metrics,
    )

    logger.info("*** Evaluate ***")

    #### eval dataset & eval example - predictions.json 생성됨
    training_args.per_device_eval_batch_size = 16
    if training_args.do_predict:
        predictions = trainer.predict(
            test_dataset=eval_dataset, test_examples=datasets["validation"]
        )

        # predictions.json 은 postprocess_qa_predictions() 호출시 이미 저장됩니다.
        logger.info(
            "No metric can be presented because there is no correct answer given. Job done!"
        )

    if training_args.do_eval:
        metrics = trainer.evaluate()
        metrics["eval_samples"] = len(eval_dataset)

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


if __name__ == "__main__":
    main()