探秘Transformers：用Hugging Face预训练模型实现命名实体识别！

Tom2Code

发布于 2024-03-02 02:07:07

1.5K00

代码可运行

文章被收录于专栏：TomTom

运行总次数：0

代码可运行

命名实体识别（NER）是自然语言处理中的一项关键任务，旨在从文本中识别并提取出具有特定意义的实体，如人名、地名、组织机构名等。通过NER，计算机可以更好地理解文本，帮助我们从海量文本数据中快速获取有用信息，是许多NLP应用的基础，如信息提取、问答系统等。

我们可以从huggingface上看一个医学实体命名的例子：

输入框里的文字是我们的输入，

点击computer按钮：

这就是我们的结果，可以这个模型成功的从我们的文本中推断出来了很多实体。例如识别出来了age，年龄。还有性别，sex。还有一些医学上的一些专业名称也成功的区分出来了，这就是命名实体识别的一个demo

那我们今天要做的demo是基于中文的命名实体识别。

首先介绍一下数据集：

可以在huggingface上直接搜索：

peoples_daily_ner

是一个报纸上的一些新闻的文字数据。

再介绍一下我们使用的预训练模型：

也是可以直接从huggingface上搜索：

hfl/chinese-macbert-base

稍微介绍一下这个模型：

MacBERT 是一种改进的 BERT，采用新颖的 MLM 作为校正预训练任务，从而减少了预训练和微调之间的差异。

在微调阶段，[MASK] 标记从未出现过，我们建议使用相似词来代替[MASK] 标记进行屏蔽。相似词通过同义词工具包（Wang 和 Hu，2017 年）获得，
该工具包基于 word2vec（Mikolov 等人，2013 年）相似性计算。
如果选择了一个 N-gram 作为掩码，我们将单独找出相似词。
在极少数情况下，如果没有相似词，我们将降级使用随机词替换。

然后就是从huggingface的镜像站上下载我们的模型：

具体的操作在上一篇我们已经介绍过：

点击图片即可进入上一篇文章

接下来废话不多说，直接开肝

1.首先是导包

import evaluate
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification

2.加载数据集

# 如果加载失败 可以通过本地下载到磁盘然后再加载
ner_datasets = load_dataset("peoples_daily_ner", cache_dir="./data")

数据集截图：

随机打印1个数据集看看：

3.加载分词器

tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-macbert-base")

还是一样，如果加载失败，可以通过先从镜像站下载到本地，然后再从本地加载分词器和模型

写一个函数，用来处理将原始文本数据中的标签（通常是实体名称或类型的标记）映射到经过标记化后的文本中的正确位置上，以便用于训练或评估模型。

def process_function(examples):
    tokenized_examples = tokenizer(examples["tokens"], max_length=128, truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_examples.word_ids(batch_index=i)
        label_ids = [-100 if word_id is None else label[word_id] for word_id in word_ids]
        labels.append(label_ids)
    tokenized_examples["labels"] = labels
    return tokenized_examples

这样数据就会变成我们想要的格式了：

4.创建模型

model = AutoModelForTokenClassification.from_pretrained("hfl/chinese-macbert-base", num_labels=len(label_list))

这里要制定我们的类别，使用list_labels来制定

5.创建评估函数

seqeval = evaluate.load("seqeval_metric.py")
seqeval

返回了很多用法

在写一个这次训练要用到的验证函数：

import numpy as np
from seqeval.metrics import f1_score

def eval_metric(pred, label_list):
    predictions, labels = pred
    predictions = np.argmax(predictions, axis=-1)

    true_predictions = [
        [label_list[p] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    true_labels = [
        [label_list[l] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    return {"f1": f1_score(true_labels, true_predictions)}

6.配置训练参数

args = TrainingArguments(
    output_dir="models_for_ner",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=64,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    metric_for_best_model="f1",
    load_best_model_at_end=True,
    logging_steps=50,
    num_train_epochs=1
)

7.创建训练器

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    compute_metrics=eval_metric,
    data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer)
)

8.模型训练

由于时间原因只训练了一个epoch

9.使用验证集来测试模型

可以看到f1值很高，从侧面也能代表模型的准确率不低。

最后就是我们再来验证一下：

from transformers import pipeline
model.config.id2label = {idx: label for idx, label in enumerate(label_list)}
ner_pipe = pipeline("token-classification", model=model, tokenizer=tokenizer, device=0, aggregation_strategy="simple")

res = ner_pipe("汤姆在北京读研究生")
res

可以看到模型准确的识别出了两个实体

可以再分割一下：

# 根据start和end取实际的结果
ner_result = {}
x = "汤姆在北京读研究生"
for r in res:
    if r["entity_group"] not in ner_result:
        ner_result[r["entity_group"]] = []
    ner_result[r["entity_group"]].append(x[r["start"]: r["end"]])

ner_result