上一篇文章讲解了Langchain,实现一个简单的demo,结合利用 LangChain 和 BERT 用于命名实体识别。
一、命名实体识别模型训练(bert+CRF)
bert作为我们的预训练模型(用于将输入文本转换为特征向量),CRF作为我们的条件随机场(将嵌入特征转为标签),既然要训练,那么我们的损失函数采用CRF 损失。
注意区分 交叉熵损失和CRF损失
CRF本身也有学习参数,一起参与梯度更新,只是参数为一块转移矩阵实现标签之间的关系建模。
实现代码如下,
模型和 分词器都是使用的bert base chinese
实现了一个结合BERT和CRF模型的命名实体识别(NER)任务。首先,定义了BertCRF
类,利用BERT进行特征提取,并通过CRF层进行序列标签预测。数据预处理部分使用BertTokenizerFast
对输入文本进行分词,同时将标签对齐到子词级别,处理特殊token。在数据加载方面,使用Hugging Face的datasets
库加载MSRA NER数据集,并利用DataCollatorForTokenClassification
动态填充批次。
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import BertTokenizerFast, BertForTokenClassification, DataCollatorForTokenClassification
from torchcrf import CRF
from torch.optim import AdamW
from datasets import load_dataset
from seqeval.metrics import classification_report, accuracy_score
from tqdm.auto import tqdm# 定义BERT + CRF模型
class BertCRF(nn.Module):def __init__(self, bert_model_name, num_labels):super(BertCRF, self).__init__()# 使用预训练的BERT模型进行特征提取self.bert = BertForTokenClassification.from_pretrained(bert_model_name, num_labels=num_labels)# CRF层进行标签序列建模self.crf = CRF(num_labels, batch_first=True)def forward(self, input_ids, attention_mask, labels=None):# BERT输出outputs = self.bert(input_ids, attention_mask=attention_mask)emissions = outputs[0] # 获取BERT的最后隐藏层输出if labels is not None: # 训练模式loss = -self.crf(emissions, labels, mask=attention_mask.bool())return losselse:predictions = self.crf.decode(emissions, mask=attention_mask.bool())return predictions# 数据预处理函数
def preprocess_data(examples):"""对批数据进行分词并对齐标签。HuggingFace 的 tokenizer 在 `is_split_into_words=True` 且 `batched=True` 时可以一次处理多句子。这里根据 `word_ids(batch_index=...)` 把原始词级别标签扩展到子词级别;对特殊 token (CLS、SEP、PAD) 使用 -100,使其在计算 loss 时被忽略。`msra_ner` 数据集的 `ner_tags` 已经是整数 ID,因此无需 label2id 转换。"""# 分词tokenized = tokenizer(examples["tokens"],