o2o-plagiarism-ai/scripts/train_klue_ner.py

"""성능지표 #3(메타 데이터 추출 F1) — KLUE NER 토큰 분류 학습·평가.

전 차수 시험(GERI `GERIR.CE.2511-02.009`) 과 동일한 과제 설정이다.

    모델   : BertForTokenClassification (klue/bert-base, num_labels=13)
    학습   : KLUE NER 약 11,000건
    평가   : 학습되지 않은 2,000여건으로 상대 평가
    결과   : micro avg f1 = 0.8113 (81%)

2-1년차 목표는 83% 다. 같은 과제를 더 큰 사전학습 모델로 올려 달성한다.
`--model klue/roberta-large` 가 기본이며, 전 차수 재현이 필요하면
`--model klue/bert-base` 로 돌린다.

출력은 전 차수 성적서와 같은 형태다. 태그별 precision/recall/f1-score/support 와
micro/macro/weighted 평균을 함께 남겨 성적서에 그대로 첨부할 수 있게 한다.

사용:
    python scripts/train_klue_ner.py --out-dir data/models/klue_ner
    python scripts/train_klue_ner.py --eval-only --out-dir data/models/klue_ner
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

# KLUE NER 태그 6종. 성적서 게재 순서이며 출력 정렬에만 쓴다.
TAGS = ["DT", "LC", "OG", "PS", "QT", "TI"]


def label_names(dataset) -> list[str]:
    """라벨 목록은 데이터셋에서 직접 가져온다. 코드에서 다시 만들지 않는다.

    KLUE NER 의 id 순서는 B-DT(0) … I-TI(11), **O(12)** 로 'O' 가 맨 뒤다.
    이를 `["O"] + [B-/I- ...]` 로 재구성하면 전체가 한 칸씩 밀려, 실제 O 구간이
    I-TI 엔티티로 집계된다. 그러면 micro support 가 14,257(실제 엔티티 수) 에서
    67,056 으로 부풀고 F1 도 함께 올라간다(0.8123 → 0.8750). 전 차수 성적서의
    support 62,760 이 이 상태였다.

    데이터셋을 유일한 출처로 삼아 그 경로를 막는다.
    """
    return list(dataset.features["ner_tags"].feature.names)


def tokenize_and_align(examples, tokenizer, max_length: int):
    """KLUE NER 은 음절 단위 라벨이라 서브워드 토큰에 맞춰 정렬한다.

    첫 서브워드에만 라벨을 주고 나머지는 -100 으로 두어 손실에서 제외한다.
    """
    encoded = tokenizer(
        examples["tokens"],
        is_split_into_words=True,
        truncation=True,
        max_length=max_length,
        padding=False,
    )
    aligned = []
    for i, tags in enumerate(examples["ner_tags"]):
        word_ids = encoded.word_ids(batch_index=i)
        previous = None
        labels = []
        for word_id in word_ids:
            if word_id is None:
                labels.append(-100)
            elif word_id != previous:
                labels.append(tags[word_id])
            else:
                labels.append(-100)
            previous = word_id
        aligned.append(labels)
    encoded["labels"] = aligned
    return encoded


def build_metrics(id2label: dict):
    from seqeval.metrics import classification_report, f1_score

    def compute(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        true_labels, true_preds = [], []
        for prediction, label in zip(predictions, labels):
            pairs = [(p, l) for p, l in zip(prediction, label) if l != -100]
            true_preds.append([id2label[int(p)] for p, _ in pairs])
            true_labels.append([id2label[int(l)] for _, l in pairs])
        report = classification_report(
            true_labels, true_preds, output_dict=True, zero_division=0)
        return {
            "f1": f1_score(true_labels, true_preds, zero_division=0),
            "report": report,
        }

    return compute


def build_scorecard(report: dict) -> dict:
    """seqeval 리포트를 전 차수 성적서와 같은 키 구성으로 재배열한다.

    성적서는 최상위에 micro 평균을 precision/recall/f1_score 로 두고,
    태그별 상세를 detailed_report 아래에 DT→LC→OG→PS→QT→TI 순으로 싣는다.
    seqeval 의 dict 순서는 태그 등장 순서라 성적서 순서와 달라지므로 여기서 고정한다.
    """
    micro = report["micro avg"]
    detailed: dict[str, dict] = {}
    for tag in TAGS:
        if tag in report:
            detailed[tag] = report[tag]
    for name in ("micro avg", "macro avg", "weighted avg"):
        if name in report:
            detailed[name] = report[name]
    return {
        "precision": micro["precision"],
        "recall": micro["recall"],
        "f1_score": micro["f1-score"],
        "detailed_report": detailed,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="klue/roberta-large",
                        help="전 차수 재현은 klue/bert-base")
    parser.add_argument("--out-dir", type=Path, default=Path("data/models/klue_ner"))
    parser.add_argument("--epochs", type=float, default=3.0)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--eval-only", action="store_true")
    args = parser.parse_args()

    import torch
    from datasets import load_dataset
    from transformers import (AutoModelForTokenClassification, AutoTokenizer,
                              DataCollatorForTokenClassification, Trainer,
                              TrainingArguments)

    dataset = load_dataset("klue", "ner")
    print("학습 %d건 / 평가 %d건"
          % (len(dataset["train"]), len(dataset["validation"])))

    labels = label_names(dataset["train"])
    label2id = {label: i for i, label in enumerate(labels)}
    id2label = {i: label for label, i in label2id.items()}
    print("라벨 %d종 (데이터셋 기준): %s" % (len(labels), ", ".join(labels)))

    source = args.out_dir if args.eval_only else args.model
    tokenizer = AutoTokenizer.from_pretrained(str(source))
    model = AutoModelForTokenClassification.from_pretrained(
        str(source), num_labels=len(labels), id2label=id2label, label2id=label2id)

    encoded = dataset.map(
        lambda batch: tokenize_and_align(batch, tokenizer, args.max_length),
        batched=True, remove_columns=dataset["train"].column_names)

    training_args = TrainingArguments(
        output_dir=str(args.out_dir / "checkpoints"),
        learning_rate=args.lr,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size * 2,
        num_train_epochs=args.epochs,
        weight_decay=0.01,
        logging_steps=100,
        save_strategy="no",
        report_to=[],
        fp16=torch.cuda.is_available(),
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=encoded["train"],
        eval_dataset=encoded["validation"],
        data_collator=DataCollatorForTokenClassification(tokenizer),
        compute_metrics=build_metrics(id2label),
    )

    if not args.eval_only:
        trainer.train()
        args.out_dir.mkdir(parents=True, exist_ok=True)
        trainer.save_model(str(args.out_dir))
        tokenizer.save_pretrained(str(args.out_dir))

    metrics = trainer.evaluate()
    report = metrics["eval_report"]
    micro = report["micro avg"]

    print("\n" + "=" * 62)
    print("메타 데이터 추출 F1-score : %.4f   [목표 0.83]" % micro["f1-score"])
    print("  precision %.4f / recall %.4f / support %d"
          % (micro["precision"], micro["recall"], micro["support"]))
    print("\n[태그별]")
    for tag in TAGS:
        if tag in report:
            row = report[tag]
            print("  %-4s P=%.4f R=%.4f F1=%.4f  support=%d"
                  % (tag, row["precision"], row["recall"],
                     row["f1-score"], row["support"]))
    for name in ("macro avg", "weighted avg"):
        row = report[name]
        print("  %-12s P=%.4f R=%.4f F1=%.4f"
              % (name, row["precision"], row["recall"], row["f1-score"]))

    # 성적서 게재 형식 그대로 출력한다. 전 차수 성적서
    # (GERI `GERIR.CE.2511-02.009`) 9페이지 "4. 결과 측정" 과 동일한 키 구성이라
    # 이 출력을 그대로 캡처해 붙일 수 있다.
    scorecard = build_scorecard(report)
    rendered = json.dumps(scorecard, ensure_ascii=False, indent=4)

    print("\n" + "=" * 62)
    print("4. 결과 측정")
    print()
    print(rendered)
    print("\n결과 : 평균 f1 -score %d%%" % round(micro["f1-score"] * 100))

    args.out_dir.mkdir(parents=True, exist_ok=True)
    (args.out_dir / "result.json").write_text(rendered + "\n", encoding="utf-8")
    (args.out_dir / "run_meta.json").write_text(
        json.dumps({"model": args.model, "epochs": args.epochs,
                    "batch_size": args.batch_size, "lr": args.lr,
                    "max_length": args.max_length},
                   ensure_ascii=False, indent=2),
        encoding="utf-8")
    print("\n%s 에 결과 기록" % (args.out_dir / "result.json"))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())