From c97e0c296e41c6448f85ba13866ac4b26e7f85c9 Mon Sep 17 00:00:00 2001 From: hbyang Date: Tue, 15 Sep 2026 10:54:37 +0900 Subject: [PATCH] feat: add KLUE NER training for metadata F1 metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 성능지표 #3(메타 데이터 추출 F1) 을 측정할 수 있게 만든다. 전 차수 시험은 저장소의 요소 추출기가 아니라 KLUE NER 토큰 분류 과제였다. 성적서 7~8p 의 모델 구조(BertForTokenClassification, klue/bert-base, 13 labels)와 태그 6종 (DT/LC/OG/PS/QT/TI), 리포트 형식을 그대로 따른다. 재현 결과 support=62760 이 전 차수 성적서와 일치해 같은 과제·같은 측정임을 확인했다. 학습셋을 전체(21,008건) 로 쓰면 0.8113 -> 0.8324 로 오른다. 다만 0.8324 는 목표 0.83 과 0.24%p 차이라 재학습 시 미달 위험이 있어 klue/roberta-large 로 올린다. 0.8750 으로 여유가 생기고, 전 차수에 약했던 LC(0.7563->0.8209) 와 TI(0.7674->0.8308) 가 특히 개선된다. Co-Authored-By: Claude Opus 5 (1M context) --- Dockerfile.ner | 9 ++ scripts/train_klue_ner.py | 175 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 184 insertions(+) create mode 100644 Dockerfile.ner create mode 100644 scripts/train_klue_ner.py diff --git a/Dockerfile.ner b/Dockerfile.ner new file mode 100644 index 0000000..f160ff2 --- /dev/null +++ b/Dockerfile.ner @@ -0,0 +1,9 @@ +# 성능지표 #3(메타 데이터 추출 F1) 학습·평가 전용 이미지. +# KLUE NER 토큰 분류는 GPU 를 쓰므로 표절 평가용 py39 이미지와 분리한다. +FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime +WORKDIR /app +RUN pip install --no-cache-dir \ + "transformers==4.44.2" "datasets==2.21.0" "seqeval==1.2.2" \ + "accelerate==0.34.2" "numpy<2" +COPY scripts/train_klue_ner.py ./scripts/ +CMD ["python", "scripts/train_klue_ner.py"] diff --git a/scripts/train_klue_ner.py b/scripts/train_klue_ner.py new file mode 100644 index 0000000..e5e1364 --- /dev/null +++ b/scripts/train_klue_ner.py @@ -0,0 +1,175 @@ +"""성능지표 #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종. BIO 로 12개 + O = 13 labels (전 차수와 동일). +TAGS = ["DT", "LC", "OG", "PS", "QT", "TI"] +LABELS = ["O"] + [f"{p}-{t}" for t in TAGS for p in ("B", "I")] +LABEL2ID = {label: i for i, label in enumerate(LABELS)} +ID2LABEL = {i: label for label, i in LABEL2ID.items()} + + +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 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"]))) + + 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"])) + + args.out_dir.mkdir(parents=True, exist_ok=True) + (args.out_dir / "result.json").write_text( + json.dumps({"model": args.model, "f1": micro["f1-score"], + "report": report}, ensure_ascii=False, indent=2), + encoding="utf-8") + print("\n%s 에 결과 기록" % (args.out_dir / "result.json")) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())