개체명 인식(NER)은 텍스트에서 사람(PER), 장소(LOC), 조직(ORG), 날짜(DAT) 등의 고유한 엔티티(Entity)를 식별하고 분류하는 태스크입니다. 텍스트 분류가 문장 전체에 하나의 레이블을 부여하는 것과 달리, NER은 각 토큰마다 레이블을 부여하는 Token Classification 태스크입니다.
BERT 토크나이저는 단어를 서브워드로 분할하므로, 원본 토큰과 서브워드 토큰 사이의 레이블 정렬이 필요합니다.
from transformers import AutoTokenizermodel_name = "klue/bert-base"tokenizer = AutoTokenizer.from_pretrained(model_name)def tokenize_and_align_labels(examples): """서브워드 토큰에 NER 레이블을 정렬하는 함수""" tokenized = tokenizer( examples["tokens"], truncation=True, is_split_into_words=True, # 이미 토큰화된 입력 max_length=128, padding="max_length", ) all_labels = [] for i, labels in enumerate(examples["ner_tags"]): word_ids = tokenized.word_ids(batch_index=i) label_ids = [] previous_word_id = None for word_id in word_ids: if word_id is None: # [CLS], [SEP], [PAD] 토큰 → -100 (손실 계산에서 무시) label_ids.append(-100) elif word_id != previous_word_id: # 새 단어의 첫 번째 서브워드 → 원본 레이블 사용 label_ids.append(labels[word_id]) else: # 같은 단어의 후속 서브워드 # B- → I- 로 변환 (선택 사항) label = labels[word_id] if label % 2 == 1: # B- 태그인 경우 label_ids.append(label + 1) # I- 태그로 변환 else: label_ids.append(label) previous_word_id = word_id all_labels.append(label_ids) tokenized["labels"] = all_labels return tokenized# 전체 데이터셋에 적용tokenized_dataset = dataset.map( tokenize_and_align_labels, batched=True, remove_columns=dataset["train"].column_names,)
레이블 정렬이 NER Fine-tuning의 핵심입니다. 서브워드 토큰에 잘못된 레이블이 할당되면 모델 성능이 크게 저하됩니다. 특수 토큰([CLS], [SEP], [PAD])에는 반드시 -100을 할당하여 손실 계산에서 제외해야 합니다.
3
모델 학습
Token Classification 모델을 로딩하고 학습합니다.
from transformers import ( AutoModelForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification,)# 모델 로딩 (각 토큰마다 13개 클래스 분류)model = AutoModelForTokenClassification.from_pretrained( model_name, num_labels=num_labels,)# Data Collator (동적 패딩)data_collator = DataCollatorForTokenClassification( tokenizer=tokenizer,)# 학습 설정training_args = TrainingArguments( output_dir="./ner-results", eval_strategy="epoch", save_strategy="epoch", learning_rate=3e-5, per_device_train_batch_size=32, per_device_eval_batch_size=64, num_train_epochs=5, weight_decay=0.01, load_best_model_at_end=True, metric_for_best_model="overall_f1", fp16=True,)
4
seqeval 기반 평가
NER은 엔티티 단위로 평가합니다. 토큰 단위 정확도가 아니라, 엔티티의 시작~끝 범위와 유형이 모두 일치해야 정답으로 인정됩니다.
import numpy as npimport evaluateseqeval = evaluate.load("seqeval")def compute_metrics(eval_pred): """seqeval 기반 NER 평가 함수""" logits, labels = eval_pred predictions = np.argmax(logits, axis=-1) # 정수 레이블을 문자열 태그로 변환 true_labels = [] true_preds = [] for pred_seq, label_seq in zip(predictions, labels): pred_tags = [] label_tags = [] for pred, label in zip(pred_seq, label_seq): if label != -100: # 특수 토큰 제외 pred_tags.append(label_names[pred]) label_tags.append(label_names[label]) true_preds.append(pred_tags) true_labels.append(label_tags) results = seqeval.compute( predictions=true_preds, references=true_labels, ) return { "precision": results["overall_precision"], "recall": results["overall_recall"], "overall_f1": results["overall_f1"], "accuracy": results["overall_accuracy"], }# Trainer 생성 및 학습trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_dataset["train"], eval_dataset=tokenized_dataset["validation"], data_collator=data_collator, compute_metrics=compute_metrics,)trainer.train()
5
추론 및 결과 시각화
학습된 모델로 새로운 텍스트에서 엔티티를 추출합니다.
from transformers import pipeline# NER 파이프라인 생성ner_pipeline = pipeline( "ner", model=trainer.model, tokenizer=tokenizer, aggregation_strategy="simple", # 서브워드 토큰을 단어 단위로 병합 device=0,)# 추론text = "삼성전자의 이재용 회장이 2024년 3월 서울에서 AI 전략 회의를 주재했다."entities = ner_pipeline(text)for entity in entities: print(f"[{entity['entity_group']}] {entity['word']} " f"(score: {entity['score']:.3f})")# [OG] 삼성전자 (score: 0.987)# [PS] 이재용 (score: 0.995)# [DT] 2024년 3월 (score: 0.972)# [LC] 서울 (score: 0.991)
seqeval은 엔티티를 하나의 단위로 평가합니다. 아래 예시를 통해 토큰 단위 평가와의 차이를 이해할 수 있습니다.
정답: B-PER I-PER O B-LOC O예측: B-PER I-PER O B-ORG O토큰 단위 정확도: 4/5 = 80%엔티티 단위 (seqeval): - "PER" 엔티티: TP (시작~끝, 유형 모두 일치) - "LOC" 엔티티: FN (놓침) - "ORG" 엔티티: FP (오탐) → Precision: 1/2, Recall: 1/2, F1: 0.50
서브워드 분할 때문에 레이블이 밀려납니다
is_split_into_words=True 파라미터를 사용하고 있는지 확인합니다. 이미 토큰화된 입력(단어 리스트)을 전달할 때 필수입니다. word_ids() 메서드로 각 서브워드가 어떤 원본 단어에 매핑되는지 추적합니다.
aggregation_strategy 옵션의 차이는 무엇인가요?
none: 서브워드 단위로 결과 반환. simple: 동일 엔티티 서브워드를 병합하고 점수 평균. first: 첫 서브워드의 점수 사용. max: 최고 점수 사용. 일반적으로 simple이 가장 직관적입니다.
CRF 레이어를 추가하면 성능이 올라가나요?
CRF(Conditional Random Field)는 태그 전이 확률을 학습하여 “B-PER 다음에 I-LOC가 오는” 비합리적 예측을 방지합니다. 데이터셋이 작을 때 효과적이지만, BERT처럼 강력한 인코더를 사용하면 CRF 없이도 충분한 성능을 달성하는 경우가 많습니다.
