Token分类任务概述

Token分类是自然语言处理中的核心序列标注任务，其本质是为文本序列中的每一个词元（Token）分配特定的类别标签。常见的应用场景包括：

• 命名实体识别（NER）：识别文本中的专有名词，如人名、地名、机构名等。
• 词性标注（POS）：为每个单词标注其词性，如名词、动词、形容词。
• 分块（Chunking）：将属于同一短语或实体的词元组合在一起，通常结合BIO（Begin, Inside, Outside）标注体系来实现。

本文将详细讲解如何基于预训练语言模型（如BERT）微调一个用于命名实体识别的Token分类模型，涵盖从数据预处理、子词标签对齐到模型训练与推理的完整流程。

数据集加载与结构探索

我们将使用经典的CoNLL-2003数据集进行演示。该数据集包含了新闻语料，并提供了NER、POS和Chunking三种标注。通过datasets库可以轻松加载并查看其结构：

from datasets import load_dataset

ner_corpus = load_dataset("conll2003")
print(ner_corpus)

输出结果会显示数据集包含训练集、验证集和测试集，且特征列中包含tokens（已预分词的单词列表）以及对应的ner_tags、pos_tags和chunk_tags。需要注意的是，这里的tokens是单词级别的，在输入模型前仍需经过Tokenizer进行子词（Subword）拆分。

在NER任务中，标签通常采用BIO格式。我们可以通过数据集的features属性来解析整数标签与字符串标签的映射关系：

ner_feature = ner_corpus["train"].features["ner_tags"]
tag_names = ner_feature.feature.names
print(tag_names)
# 输出: ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC']

其中，O代表非实体，B-前缀表示实体的起始词，I-前缀表示实体的内部词。

子词对齐与数据预处理

由于预训练模型使用的Tokenizer会将单词拆分为更小的子词（例如将"lamb"拆分为"la"和"##mb"），这会导致输入Token的数量与原始单词标签的数量不一致。此外，Tokenizer还会添加[CLS]和[SEP]等特殊Token。因此，我们必须将原始单词级别的标签对齐到子词级别。

对齐的核心规则如下：

1. 特殊Token（如[CLS]、[SEP]）的标签设置为-100，以便在计算交叉熵损失时被自动忽略。
2. 单词拆分后的第一个子词保留原始单词的标签。
3. 单词拆分后的后续子词，如果原标签是B-开头，则需转换为I-开头，以保证实体边界的逻辑正确性。

以下是实现标签对齐的函数：

def map_tags_to_subwords(original_tags, word_id_mapping):
    aligned_tags = []
    prev_word_idx = None
    for current_word_idx in word_id_mapping:
        if current_word_idx is None:
            # 特殊Token标记为-100
            aligned_tags.append(-100)
        elif current_word_idx != prev_word_idx:
            # 单词的第一个子词，保留原标签
            aligned_tags.append(original_tags[current_word_idx])
        else:
            # 单词的后续子词
            tag = original_tags[current_word_idx]
            # 在标准BIO编码中，B-标签通常为奇数，I-标签为偶数
            aligned_tags.append(tag + 1 if tag % 2 != 0 else tag)
        prev_word_idx = current_word_idx
    return aligned_tags

接下来，我们将Tokenizer与对齐逻辑结合，编写批量处理函数，并利用Dataset.map对整个数据集进行预处理：

from transformers import AutoTokenizer

model_id = "bert-base-cased"
text_tokenizer = AutoTokenizer.from_pretrained(model_id)

def process_batch(batch_data):
    tokenized = text_tokenizer(
        batch_data["tokens"], 
        truncation=True, 
        is_split_into_words=True
    )
    updated_labels = []
    for idx, tags in enumerate(batch_data["ner_tags"]):
        w_ids = tokenized.word_ids(batch_index=idx)
        updated_labels.append(map_tags_to_subwords(tags, w_ids))
    
    tokenized["labels"] = updated_labels
    return tokenized

tokenized_corpus = ner_corpus.map(
    process_batch,
    batched=True,
    remove_columns=ner_corpus["train"].column_names,
)

数据整理器与评估指标

在构建Batch时，除了对输入ID进行填充，还需要对标签进行同步填充。我们可以使用DataCollatorForTokenClassification，它会自动使用-100对标签进行填充：

from transformers import DataCollatorForTokenClassification

data_collator = DataCollatorForTokenClassification(tokenizer=text_tokenizer)

对于Token分类任务，标准的评估库是seqeval。由于该库要求输入字符串格式的标签，我们需要在评估函数中将模型输出的Logits转换为字符串标签，并过滤掉-100的占位符：

import evaluate
import numpy as np

seqeval_metric = evaluate.load("seqeval")

def evaluate_predictions(eval_data):
    logits, true_labels = eval_data
    preds = np.argmax(logits, axis=-1)
    
    # 过滤-100并转换为字符串标签
    pred_strings = [[tag_names[p] for p, l in zip(pred, lbl) if l != -100] for pred, lbl in zip(preds, true_labels)]
    label_strings = [[tag_names[l] for l in lbl if l != -100] for lbl in true_labels]
    
    metrics = seqeval_metric.compute(predictions=pred_strings, references=label_strings)
    return {
        "f1_score": metrics["overall_f1"],
        "precision": metrics["overall_precision"],
        "recall": metrics["overall_recall"],
        "accuracy": metrics["overall_accuracy"]
    }

模型初始化与Trainer微调

使用AutoModelForTokenClassification加载预训练模型。为了确保模型在Hub上展示正确的标签映射，我们需要显式传入id2label和label2id字典：

from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer

id2label = {str(i): label for i, label in enumerate(tag_names)}
label2id = {v: k for k, v in id2label.items()}

model = AutoModelForTokenClassification.from_pretrained(
    model_id,
    id2label=id2label,
    label2id=label2id,
)

training_args = TrainingArguments(
    output_dir="my-ner-model",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    weight_decay=0.01,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_corpus["train"],
    eval_dataset=tokenized_corpus["validation"],
    data_collator=data_collator,
    compute_metrics=evaluate_predictions,
    tokenizer=text_tokenizer,
)

trainer.train()
trainer.push_to_hub(commit_message="Training completed")

基于Accelerate的自定义训练循环

如果需要更细粒度的控制，可以使用Accelerate库构建原生PyTorch训练循环。以下是分布式环境下的完整训练逻辑：

from torch.utils.data import DataLoader
from torch.optim import AdamW
from transformers import get_scheduler
from accelerate import Accelerator
from tqdm.auto import tqdm
import torch

# 初始化Accelerator
accelerator = Accelerator()

train_loader = DataLoader(tokenized_corpus["train"], shuffle=True, collate_fn=data_collator, batch_size=8)
eval_loader = DataLoader(tokenized_corpus["validation"], collate_fn=data_collator, batch_size=8)

optimizer = AdamW(model.parameters(), lr=2e-5)

num_epochs = 3
num_training_steps = num_epochs * len(train_loader)
lr_scheduler = get_scheduler("linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps)

model, optimizer, train_loader, eval_loader = accelerator.prepare(model, optimizer, train_loader, eval_loader)

training_progress = tqdm(total=num_training_steps)

for current_epoch in range(num_epochs):
    # 训练阶段
    model.train()
    for batch in train_loader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        training_progress.update(1)

    # 评估阶段
    model.eval()
    for batch in eval_loader:
        with torch.no_grad():
            outputs = model(**batch)
        
        preds = outputs.logits.argmax(dim=-1)
        targets = batch["labels"]
        
        # 跨进程填充以对齐形状，防止gather时维度不匹配
        preds = accelerator.pad_across_processes(preds, dim=1, pad_index=-100)
        targets = accelerator.pad_across_processes(targets, dim=1, pad_index=-100)
        
        gathered_preds = accelerator.gather(preds)
        gathered_targets = accelerator.gather(targets)
        
        # 此处需调用前文定义的格式转换逻辑将gathered结果转为seqeval所需格式
        # final_preds, final_targets = format_for_seqeval(gathered_preds, gathered_targets)
        # seqeval_metric.add_batch(predictions=final_preds, references=final_targets)

    # 保存模型
    accelerator.wait_for_everyone()
    base_model = accelerator.unwrap_model(model)
    base_model.save_pretrained("my-ner-model", save_function=accelerator.save)
    if accelerator.is_main_process:
        text_tokenizer.save_pretrained("my-ner-model")

模型推理与应用

微调完成后，可以通过pipeline快速加载模型进行推理。设置aggregation_strategy="simple"可以将子词级别的预测结果自动聚合成完整的实体词汇：

from transformers import pipeline

ner_pipeline = pipeline(
    "token-classification", 
    model="my-ner-model", 
    aggregation_strategy="simple"
)

text = "Apple is looking at buying U.K. startup for $1 billion"
results = ner_pipeline(text)
print(results)

输出结果将直接返回识别出的实体组、置信度分数以及在原句中的字符起止位置，方便直接集成到下游业务系统中。