Обучение модели NER (Named Entity Recognition)
Fine-tuning NER-модели на кастомных типах сущностей — стандартная задача при работе с доменными текстами: медицинскими, юридическими, финансовыми документами.
Подготовка данных: IOB2-разметка
Стандартный формат для NER — IOB2 (BIO-tagging):
-
B-ORG— начало сущности типа ORG -
I-ORG— продолжение сущности ORG -
O— не сущность
Газпром B-ORG
подписал O
контракт O
с O
Deutsche B-ORG
Bank I-ORG
в O
Берлине B-LOC
Разметка: Prodigy (дорого, удобно), Label Studio (open-source), Doccano (простой UI). Минимум 200 предложений на тип сущности для приемлемого качества, 1000+ для высокого.
Fine-tuning на TokenClassification
from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer
from transformers import DataCollatorForTokenClassification
label_list = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
label2id = {l: i for i, l in enumerate(label_list)}
id2label = {i: l for l, i in label2id.items()}
model = AutoModelForTokenClassification.from_pretrained(
"DeepPavlov/rubert-base-cased",
num_labels=len(label_list),
id2label=id2label,
label2id=label2id
)
data_collator = DataCollatorForTokenClassification(tokenizer)
training_args = TrainingArguments(
output_dir="./ner_model",
num_train_epochs=10, # NER требует больше эпох чем классификация
per_device_train_batch_size=16,
learning_rate=5e-5,
weight_decay=0.01,
)
Метрики NER: seqeval
import evaluate
seqeval = evaluate.load("seqeval")
def compute_metrics(p):
predictions, labels = p
predictions = np.argmax(predictions, axis=2)
true_predictions = [
[label_list[p] for (p, l) in zip(pred, label) if l != -100]
for pred, label in zip(predictions, labels)
]
true_labels = [
[label_list[l] for (p, l) in zip(pred, label) if l != -100]
for pred, label in zip(predictions, labels)
]
results = seqeval.compute(predictions=true_predictions, references=true_labels)
return {"f1": results["overall_f1"], "precision": results["overall_precision"]}
Работа с несбалансированными типами
Редкие типы сущностей обучаются плохо. Стратегии: oversampling предложений с редкими сущностями, аугментация (замена имён из словаря при сохранении типа), упрощение схемы (объединить редкие типы).
Типичные F1 после fine-tuning: PER 94–97%, ORG 88–93%, кастомные доменные сущности 80–92%.