实战篇----利用 LangChain 和 BERT 用于命名实体识别-----完整代码

上一篇文章讲解了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 loss
        else:
            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"],