train_squad.py

"""
SQuAD v1.1에 BERT를 파인튜닝하는 스크립트
==========================================
Devlin et al. (2019)의 §4.2 SQuAD 실험을 재현합니다.

논문은 추출형 QA를 passage 토큰들에 대한 두 개의 확률 분포 예측 문제로
정의합니다 - 답변 span의 시작 위치와 끝 위치 각각:

    P_i = exp(S · T_i) / Σ_j exp(S · T_j)

끝 토큰도 같은 형태로, 학습되는 끝 벡터 E를 사용합니다. span (i, j)의
점수는 S·T_i + E·T_j 이고, j >= i 조건에서 가장 높은 점수의 span을 반환합니다.

§4.2의 하이퍼파라미터:
    에폭 수    : 3
    학습률     : 5e-5
    배치 크기  : 32

실행 예시
---------
    python train_squad.py \\
        --model_name_or_path bert-base-uncased \\
        --output_dir ./out/squad \\
        --learning_rate 5e-5 \\
        --num_train_epochs 3 \\
        --per_device_train_batch_size 32
"""

from __future__ import annotations

import argparse

from datasets import load_dataset
from transformers import (
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DefaultDataCollator,
    Trainer,
    TrainingArguments,
    set_seed,
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--model_name_or_path", type=str, default="bert-base-uncased")
    p.add_argument("--output_dir", type=str, default="./out/squad")
    p.add_argument("--max_length", type=int, default=384,
                   help="시퀀스 최대 길이. 논문은 SQuAD에서 384를 사용.")
    p.add_argument("--doc_stride", type=int, default=128,
                   help="긴 지문을 겹치는 윈도우로 나눌 때의 stride.")
    p.add_argument("--learning_rate", type=float, default=5e-5,
                   help="논문: 5e-5 (§4.2).")
    p.add_argument("--num_train_epochs", type=float, default=3.0,
                   help="논문: 3 에폭 (§4.2).")
    p.add_argument("--per_device_train_batch_size", type=int, default=32,
                   help="논문: 배치 크기 32 (§4.2).")
    p.add_argument("--per_device_eval_batch_size", type=int, default=32)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def main() -> None:
    args = parse_args()
    set_seed(args.seed)

    # ------------------------------------------------------------------
    # 1. 데이터셋 로드
    # ------------------------------------------------------------------
    raw = load_dataset("squad")

    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)

    def preprocess_train(examples):
        # 논문은 (질문, 지문)을 단일 시퀀스로 패킹합니다. 질문은 segment 임베딩 A,
        # 지문은 segment 임베딩 B를 사용합니다.
        # `tokenizer(question, context, ...)`가 [SEP]와 segment id를 처리해줍니다.
        questions = [q.strip() for q in examples["question"]]
        tokenized = tokenizer(
            questions,
            examples["context"],
            max_length=args.max_length,
            truncation="only_second",     # 질문은 절대 자르지 않음
            stride=args.doc_stride,        # 긴 지문은 겹치는 윈도우로 분할
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length",
        )

        # 하나의 지문이 여러 feature로 나뉘면, 각 feature를 원본 예제로 매핑하여
        # 답변 span을 찾을 수 있게 합니다.
        sample_mapping = tokenized.pop("overflow_to_sample_mapping")
        offset_mapping = tokenized.pop("offset_mapping")

        tokenized["start_positions"] = []
        tokenized["end_positions"] = []

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

            sequence_ids = tokenized.sequence_ids(i)
            sample_idx = sample_mapping[i]
            answers = examples["answers"][sample_idx]

            if len(answers["answer_start"]) == 0:
                # §4.3: SQuAD v1.1은 모든 질문에 답이 있습니다.
                # 만약 v2.0을 돌린다면 답이 없는 경우 [CLS]를 가리키게 합니다.
                tokenized["start_positions"].append(cls_index)
                tokenized["end_positions"].append(cls_index)
                continue

            start_char = answers["answer_start"][0]
            end_char = start_char + len(answers["text"][0])

            # 지문 토큰 범위를 찾습니다 (sequence_id == 1).
            token_start_index = 0
            while sequence_ids[token_start_index] != 1:
                token_start_index += 1
            token_end_index = len(input_ids) - 1
            while sequence_ids[token_end_index] != 1:
                token_end_index -= 1

            # 답이 현재 윈도우 밖에 있으면 [CLS]를 가리킵니다.
            if not (
                offsets[token_start_index][0] <= start_char
                and offsets[token_end_index][1] >= end_char
            ):
                tokenized["start_positions"].append(cls_index)
                tokenized["end_positions"].append(cls_index)
            else:
                while (
                    token_start_index < len(offsets)
                    and offsets[token_start_index][0] <= start_char
                ):
                    token_start_index += 1
                tokenized["start_positions"].append(token_start_index - 1)
                while offsets[token_end_index][1] >= end_char:
                    token_end_index -= 1
                tokenized["end_positions"].append(token_end_index + 1)

        return tokenized

    train_dataset = raw["train"].map(
        preprocess_train,
        batched=True,
        remove_columns=raw["train"].column_names,
        desc="train 셋 토큰화 중",
    )

    # 평가에는 offset과 example_id를 유지해야 모델 출력을 문자 span으로
    # 다시 매핑할 수 있습니다 (post-processing은 여기선 생략 -
    # HuggingFace의 run_qa.py 참고).
    eval_dataset = raw["validation"].map(
        preprocess_train,
        batched=True,
        remove_columns=raw["validation"].column_names,
        desc="validation 셋 토큰화 중",
    )

    # ------------------------------------------------------------------
    # 2. 모델
    # ------------------------------------------------------------------
    model = AutoModelForQuestionAnswering.from_pretrained(args.model_name_or_path)

    # ------------------------------------------------------------------
    # 3. 학습
    # ------------------------------------------------------------------
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        learning_rate=args.learning_rate,
        num_train_epochs=args.num_train_epochs,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        weight_decay=0.01,
        warmup_ratio=0.1,
        seed=args.seed,
        report_to="none",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        data_collator=DefaultDataCollator(),
    )

    trainer.train()
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    print(f"\n모델이 {args.output_dir}에 저장되었습니다.")
    print("참고: 공식 SQuAD EM/F1 메트릭을 얻으려면 post-processing이 필요합니다")
    print("(https://github.com/huggingface/transformers/tree/main/examples/pytorch/question-answering 참고).")


if __name__ == "__main__":
    main()