如何基于B-I-O标注方案微调BERT实现自定义NER任务？

一、按B-I-O规则标注数据 #

B-I-O是NER任务的标准标注体系：

• O：非实体token
• B-XXX：XXX类型实体的起始token
• I-XXX：XXX类型实体的后续token

针对你的示例文本，标注后的数据格式如下（每个token对应一个标签）：

将标注数据整理为DataFrame：

import pandas as pd

annotated_data = [
    (["The", "income", "statement", "shows"], ["O", "B-INC", "I-INC", "O"]),
    (["Total", "expenses", "for", "the", "quarter", "are"], ["O", "B-EXP", "O", "O", "O", "O"]),
    (["Net", "income", "is", "positive."], ["O", "B-INC", "O", "O"])
]

df_annotated = pd.DataFrame(annotated_data, columns=["tokens", "labels"])

如果数据量较大，可使用LabelStudio等标注工具批量完成标注。

二、预处理数据适配BERT输入 #

BERT会对长token做subword分割，需确保标签与分割后的token对齐，忽略特殊token（如[CLS]、[SEP]）的损失计算：

from transformers import BertTokenizer
from datasets import Dataset

tokenizer = BertTokenizer.from_pretrained("dslim/bert-base-NER")

def preprocess_function(examples):
    encoding = tokenizer(examples["tokens"], is_split_into_words=True, truncation=True, padding="max_length")
    label_ids = []
    
    for idx, label in enumerate(examples["labels"]):
        word_ids = encoding.word_ids(batch_index=idx)
        previous_word_idx = None
        current_label_ids = []
        
        for word_idx in word_ids:
            if word_idx is None:
                # 特殊token对应标签设为-100，模型训练时会忽略
                current_label_ids.append(-100)
            elif word_idx != previous_word_idx:
                # 新token使用原始标签
                current_label_ids.append(label[word_idx])
            else:
                # 同一实体的subtoken，将B-开头转为I-
                if label[word_idx].startswith("B-"):
                    current_label_ids.append(label[word_idx].replace("B-", "I-"))
                else:
                    current_label_ids.append(label[word_idx])
            previous_word_idx = word_idx
        
        label_ids.append(current_label_ids)
    
    encoding["labels"] = label_ids
    return encoding

# 转换为Dataset格式
dataset = Dataset.from_pandas(df_annotated)
encoded_dataset = dataset.map(preprocess_function, batched=True)

三、微调BERT-NER模型 #

加载预训练模型，替换分类头以适配自定义标签，然后启动训练：

from transformers import BertForTokenClassification, TrainingArguments, Trainer
import evaluate
import numpy as np

# 定义自定义标签列表
label_list = ["O", "B-INC", "I-INC", "B-EXP", "I-EXP"]
num_labels = len(label_list)

# 加载预训练模型并修改分类头
model = BertForTokenClassification.from_pretrained(
    "dslim/bert-base-NER",
    num_labels=num_labels,
    id2label={i: label for i, label in enumerate(label_list)},
    label2id={label: i for i, label in enumerate(label_list)}
)

# 设置训练参数
training_args = TrainingArguments(
    output_dir="./bert-custom-ner",
    per_device_train_batch_size=8,
    num_train_epochs=3,
    logging_dir="./logs",
    logging_steps=10,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True
)

# 定义指标计算函数
metric = evaluate.load("seqeval")

def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    
    # 过滤-100标签，匹配真实标签和预测标签
    true_predictions = [
        [label_list[p] for (p, l) in zip(pred, lab) if l != -100]
        for pred, lab in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(pred, lab) if l != -100]
        for pred, lab in zip(predictions, labels)
    ]
    
    results = metric.compute(predictions=true_predictions, references=true_labels)
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"]
    }

# 初始化训练器并启动训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=encoded_dataset, # 实际需拆分训练/验证集
    eval_dataset=encoded_dataset,
    compute_metrics=compute_metrics
)

trainer.train()

注意：实际训练时需将数据集拆分为训练集和验证集，避免过拟合。

四、使用微调后的模型预测 #

加载训练好的最优模型，进行实体识别：

from transformers import pipeline

fine_tuned_ner = pipeline(
    "ner",
    model="./bert-custom-ner/checkpoint-best",
    tokenizer=tokenizer,
    aggregation_strategy="simple" # 合并同一实体的连续token
)

# 测试原始文本
test_df = pd.DataFrame({'text': ["The income statement shows", "Total expenses for the quarter are", "Net income is positive."]})
test_df['entities'] = test_df['text'].apply(lambda text: fine_tuned_ner(text))
print(test_df['entities'])

此时模型会识别出INC和EXP类型的实体，不再返回空列表。

内容的提问来源于stack exchange，提问作者kms