diff --git a/scripts/show_klue_ner_samples.py b/scripts/show_klue_ner_samples.py index a408651..0349ad7 100644 --- a/scripts/show_klue_ner_samples.py +++ b/scripts/show_klue_ner_samples.py @@ -52,6 +52,69 @@ def to_inline(tokens: list[str], tags: list[str]) -> str: 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")) @@ -59,6 +122,12 @@ def main() -> int: 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 @@ -67,6 +136,14 @@ def main() -> int: 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: @@ -86,10 +163,15 @@ def main() -> int: lines.append("") row = ds[i] sentence = to_inline(row["tokens"], [names[t] for t in row["ner_tags"]]) - lines.append('"id": %d,' % i) + 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: