o2o-plagiarism-ai/scripts/show_klue_ner_samples.py
hbyang 09647a3958 feat: add id offset and prediction array output to NER sample viewer
성적서 '3. 테스트 데이터 예시' 는 평가 데이터 id 를 학습 데이터 뒤에 이어서
매기고(11480~), 바로 아래 '예측 결과 예시' 에 마지막 예시 문장의 태그 배열을
전체로 싣는다. 두 가지를 스크립트에서 바로 낼 수 있게 한다.

--id-offset : id 시작값. 미지정 시 validation 은 학습셋 크기만큼 이어서 매김
--predict   : 마지막 예시 문장의 original_tags / predicted_tags 를 전체 배열로 출력

공백 토큰은 서브워드가 없어 word_ids 에서 빠지므로, 음절 수만큼 자리를 잡아
두고 word_id 위치에만 예측을 채워 정답과 길이를 맞춘다.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 14:40:29 +09:00

187 lines
7.0 KiB
Python

"""성능지표 #3 학습/평가 데이터 예시를 성적서 형식으로 출력한다.
전 차수 성적서(GERI `GERIR.CE.2511-02.009`) 8페이지 "1. 데이터 준비" 는 KLUE NER 을
원본 배포 형식인 인라인 태그로 싣는다.
"id": 0,
"sentence": "특히 <영동고속도로:LC> <강릉:LC> 방향 <문막휴게소:LC>에서 ..."
HuggingFace `datasets` 의 klue/ner 는 음절 토큰(`tokens`)과 BIO 태그(`ner_tags`)
배열로 들어오므로 모양이 다르다. 여기서 BIO 를 엔티티 단위로 묶어 원본 형식으로
되돌린다. 성적서에 그대로 캡처해 붙이기 위한 열람 전용 스크립트이며 학습에는
관여하지 않는다.
사용:
python scripts/show_klue_ner_samples.py # 학습셋 앞뒤 3건씩
python scripts/show_klue_ner_samples.py --split validation
python scripts/show_klue_ner_samples.py --head 5 --tail 5
python scripts/show_klue_ner_samples.py --all --out train_samples.txt
"""
from __future__ import annotations
import argparse
import sys
def to_inline(tokens: list[str], tags: list[str]) -> str:
"""음절 토큰 + BIO 태그 → `<표면형:태그>` 인라인 문장.
KLUE NER 의 tokens 는 공백을 포함한 음절 단위라 그대로 이어 붙이면 원문이 된다.
"""
out: list[str] = []
buf: list[str] = []
label: str | None = None
def flush() -> None:
if buf:
out.append("<%s:%s>" % ("".join(buf), label))
buf.clear()
for token, tag in zip(tokens, tags):
if tag.startswith("B-"):
flush()
label = tag[2:]
buf.append(token)
elif tag.startswith("I-") and label == tag[2:]:
buf.append(token)
else:
flush()
label = None
out.append(token)
flush()
return "".join(out)
def _wrap(prefix: str, tags: list[str], width: int = 88) -> list[str]:
"""성적서 지면 폭에 맞춰 배열을 줄바꿈한다."""
items = ['"%s"' % t for t in tags]
out, line = [], prefix + "["
for n, item in enumerate(items):
piece = item + (", " if n < len(items) - 1 else "]")
if len(line) + len(piece) > width and line.strip() not in ("[", prefix + "["):
out.append(line)
line = ""
line += piece
out.append(line)
return out
def predict_block(ds, names: list[str], index: int, model_dir: str) -> list[str]:
"""한 문장의 정답 태그와 예측 태그를 전체 배열로 낸다.
토크나이저가 공백 토큰에는 서브워드를 만들지 않아 word_ids 에서 빠진다.
그대로 이어 붙이면 정답과 길이가 어긋나므로, 음절 수만큼 자리를 잡아 두고
word_id 위치에만 예측을 채운다(공백은 엔티티가 될 수 없어 O 로 남는다).
"""
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForTokenClassification.from_pretrained(model_dir).eval()
if torch.cuda.is_available():
model = model.cuda()
row = ds[index]
tokens = row["tokens"]
encoded = tokenizer(tokens, is_split_into_words=True, truncation=True,
max_length=256, return_tensors="pt")
if torch.cuda.is_available():
encoded = {k: v.cuda() for k, v in encoded.items()}
word_ids = tokenizer(tokens, is_split_into_words=True, truncation=True,
max_length=256).word_ids()
else:
word_ids = encoded.word_ids(0)
encoded = dict(encoded)
with torch.no_grad():
predicted_ids = model(**encoded).logits.argmax(-1)[0].tolist()
predicted = ["O"] * len(tokens)
previous = None
for position, word_id in enumerate(word_ids):
if word_id is None or word_id == previous:
previous = word_id
continue
previous = word_id
predicted[word_id] = model.config.id2label[predicted_ids[position]]
original = [names[t] for t in row["ner_tags"]]
matched = sum(1 for a, b in zip(original, predicted) if a == b)
lines = _wrap('"original_tags": ', original)
lines.append("")
lines.extend(_wrap('"predicted_tags": ', predicted))
lines.append("")
lines.append("음절 일치 %d / %d" % (matched, len(original)))
return lines
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--split", default="train", choices=("train", "validation"))
parser.add_argument("--head", type=int, default=3, help="앞에서 N건")
parser.add_argument("--tail", type=int, default=3, help="뒤에서 N건")
parser.add_argument("--all", action="store_true", help="전건 출력")
parser.add_argument("--out", help="파일로 저장 (미지정 시 표준출력)")
parser.add_argument("--id-offset", type=int, default=None,
help="id 시작값. 미지정 시 validation 은 학습셋 크기만큼 "
"이어서 매긴다(성적서 표기 방식).")
parser.add_argument("--predict", metavar="MODEL_DIR",
help="마지막 예시 문장의 original_tags / predicted_tags 를 "
"전체 배열로 출력 (성적서 '예측 결과 예시')")
args = parser.parse_args()
from datasets import load_dataset
ds = load_dataset("klue", "ner", split=args.split)
names = ds.features["ner_tags"].feature.names
total = len(ds)
# 성적서는 평가 데이터 id 를 학습 데이터 뒤에 이어서 매긴다(전 차수 8~9p).
if args.id_offset is not None:
offset = args.id_offset
elif args.split == "validation":
offset = len(load_dataset("klue", "ner", split="train"))
else:
offset = 0
if args.all:
indices = list(range(total))
else:
indices = list(range(min(args.head, total)))
tail_start = max(len(indices), total - args.tail)
indices += list(range(tail_start, total))
lines: list[str] = []
kind = "학습" if args.split == "train" else "평가"
lines.append("<%s 데이터 예시(총 %d 건)>" % (kind, total))
lines.append("")
previous = -1
for i in indices:
if previous >= 0 and i != previous + 1:
lines.append("...")
lines.append("")
row = ds[i]
sentence = to_inline(row["tokens"], [names[t] for t in row["ner_tags"]])
lines.append('"id": %d,' % (i + offset))
lines.append('"sentence": "%s",' % sentence)
previous = i
if args.predict:
lines.append("")
lines.append("예측 결과 예시")
lines.extend(predict_block(ds, names, indices[-1], args.predict))
text = "\n".join(lines)
if args.out:
with open(args.out, "w", encoding="utf-8") as f:
f.write(text + "\n")
print("%s 에 %d건 기록 (전체 %d건)" % (args.out, len(indices), total))
else:
print(text)
return 0
if __name__ == "__main__":
raise SystemExit(main())