当前位置： 首页 > article >正文

基于BERT的命名体识别（NER）

article 2024/11/17 6:15:22

基于BERT的命名实体识别（NER）

目录

  1. 项目背景
  2. 项目结构
  3. 环境准备
  4. 数据准备
  5. 代码实现
    • 5.1 数据预处理 (src/preprocess.py)
    • 5.2 模型训练 (src/train.py)
    • 5.3 模型评估 (src/evaluate.py)
    • 5.4 模型推理 (src/inference.py)
  6. 项目运行
    • 6.1 一键运行脚本 (run.sh)
    • 6.2 手动运行
  7. 结果展示
  8. 结论
  9. 参考资料

1. 项目背景

命名实体识别（Named Entity Recognition，NER）是自然语言处理（NLP）中的基础任务之一，旨在从非结构化文本中自动识别并分类出具有特定意义的实体，例如人名、地名、组织机构名等。随着预训练语言模型（如BERT）的出现，NER的性能得到了显著提升。本项目基于BERT模型，完成对文本的序列标注，实现命名实体识别。

2. 项目结构

bert-ner/
├── data/
│   ├── train.txt            # 训练数据
│   ├── dev.txt              # 验证数据
│   ├── label_list.txt       # 标签列表
├── src/
│   ├── preprocess.py        # 数据预处理模块
│   ├── train.py             # 模型训练脚本
│   ├── evaluate.py          # 模型评估脚本
│   ├── inference.py         # 模型推理脚本
├── models/
│   ├── bert_ner_model/      # 训练好的模型文件夹
│       ├── config.json      # 模型配置文件
│       ├── pytorch_model.bin# 模型权重
│       ├── vocab.txt        # 词汇表
│       ├── tokenizer.json   # 分词器配置
│       ├── label2id.json    # 标签到ID的映射
│       ├── id2label.json    # ID到标签的映射
├── README.md                # 项目说明文档
├── requirements.txt         # 项目依赖包列表
└── run.sh                   # 一键运行脚本

3. 环境准备

3.1 创建虚拟环境（可选）

建议使用Python虚拟环境来隔离项目依赖，防止版本冲突。

# 创建虚拟环境
python -m venv venv

# 激活虚拟环境（Linux/MacOS）
source venv/bin/activate

# 激活虚拟环境（Windows）
venv\Scripts\activate

3.2 安装依赖

使用requirements.txt安装项目所需的依赖包。

pip install -r requirements.txt

requirements.txt内容：

torch==1.11.0
transformers==4.18.0
seqeval==1.2.2

注意：请根据您的Python版本和环境，选择合适的torch版本。

4. 数据准备

4.1 数据格式

训练和验证数据应采用以下格式，每行包含一个单词及其对应的标签，空行表示一个句子的结束：

John B-PER
lives O
in O
New B-LOC
York I-LOC
City I-LOC
. O

He O
works O
at O
Google B-ORG
. O

4.2 标签列表

创建label_list.txt文件，包含所有可能的标签，每行一个标签，例如：

O
B-PER
I-PER
B-ORG
I-ORG
B-LOC
I-LOC
B-MISC
I-MISC

5. 代码实现

5.1 数据预处理 (src/preprocess.py)

import torch
from torch.utils.data import Dataset
from transformers import BertTokenizer

class NERDataset(Dataset):
    """
    自定义Dataset类，用于加载NER数据。
    """
    def __init__(self, data_path, tokenizer, label2id, max_len=128):
        """
        初始化函数。

        Args:
            data_path (str): 数据文件路径。
            tokenizer (BertTokenizer): BERT分词器。
            label2id (dict): 标签到ID的映射。
            max_len (int): 序列最大长度。
        """
        self.tokenizer = tokenizer
        self.label2id = label2id
        self.max_len = max_len
        self.texts, self.labels = self._read_data(data_path)

    def _read_data(self, path):
        """
        读取数据文件。

        Args:
            path (str): 数据文件路径。

        Returns:
            texts (List[List[str]]): 文本序列列表。
            labels (List[List[str]]): 标签序列列表。
        """
        texts, labels = [], []
        with open(path, 'r', encoding='utf-8') as f:
            words, tags = [], []
            for line in f:
                if line.strip() == '':
                    if words:
                        texts.append(words)
                        labels.append(tags)
                        words, tags = [], []
                else:
                    splits = line.strip().split()
                    if len(splits) != 2:
                        continue
                    word, tag = splits
                    words.append(word)
                    tags.append(tag)
            if words:
                texts.append(words)
                labels.append(tags)
        return texts, labels

    def __len__(self):
        """
        返回数据集大小。

        Returns:
            int: 数据集大小。
        """
        return len(self.texts)

    def __getitem__(self, idx):
        """
        获取指定索引的数据样本。

        Args:
            idx (int): 索引。

        Returns:
            dict: 包含input_ids、attention_mask、labels的字典。
        """
        words, labels = self.texts[idx], self.labels[idx]
        encoding = self.tokenizer(
            words,
            is_split_into_words=True,
            return_offsets_mapping=True,
            padding='max_length',
            truncation=True,
            max_length=self.max_len
        )
        offset_mappings = encoding.pop('offset_mapping')
        labels_ids = []
        for idx, word_id in enumerate(encoding.word_ids()):
            if word_id is None:
                labels_ids.append(-100)  # 忽略[CLS], [SEP]等特殊标记
            else:
                labels_ids.append(self.label2id.get(labels[word_id], self.label2id['O']))
        encoding['labels'] = labels_ids
        # 将所有值转换为tensor
        return {key: torch.tensor(val) for key, val in encoding.items()}

5.2 模型训练 (src/train.py)

import argparse
import os
import json
import torch
from torch.utils.data import DataLoader
from transformers import BertForTokenClassification, BertTokenizer, AdamW, get_linear_schedule_with_warmup
from preprocess import NERDataset

def load_labels(label_path):
    """
    加载标签列表，并创建标签与ID之间的映射。

    Args:
        label_path (str): 标签列表文件路径。

    Returns:
        labels (List[str]): 标签列表。
        label2id (dict): 标签到ID的映射。
        id2label (dict): ID到标签的映射。
    """
    with open(label_path, 'r', encoding='utf-8') as f:
        labels = [line.strip() for line in f]
    label2id = {label: idx for idx, label in enumerate(labels)}
    id2label = {idx: label for idx, label in enumerate(labels)}
    return labels, label2id, id2label

def train(args):
    """
    模型训练主函数。

    Args:
        args (argparse.Namespace): 命令行参数。
    """
    # 加载标签和分词器
    labels, label2id, id2label = load_labels(args.label_list)
    tokenizer = BertTokenizer.from_pretrained(args.pretrained_model)
    model = BertForTokenClassification.from_pretrained(
        args.pretrained_model, num_labels=len(labels)
    )

    # 加载训练数据
    train_dataset = NERDataset(args.train_data, tokenizer, label2id, args.max_len)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)

    # 设置优化器和学习率调度器
    optimizer = AdamW(model.parameters(), lr=args.lr)
    total_steps = len(train_loader) * args.epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=int(0.1 * total_steps), num_training_steps=total_steps
    )

    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    # 创建模型保存目录
    if not os.path.exists(args.model_dir):
        os.makedirs(args.model_dir)

    # 模型训练
    model.train()
    for epoch in range(args.epochs):
        total_loss = 0
        for batch in train_loader:
            optimizer.zero_grad()
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            scheduler.step()
            total_loss += loss.item()
        avg_loss = total_loss / len(train_loader)
        print(f'Epoch {epoch+1}/{args.epochs}, Loss: {avg_loss:.4f}')
    
    # 保存模型和分词器
    model.save_pretrained(args.model_dir)
    tokenizer.save_pretrained(args.model_dir)
    # 保存标签映射
    with open(os.path.join(args.model_dir, 'label2id.json'), 'w') as f:
        json.dump(label2id, f)
    with open(os.path.join(args.model_dir, 'id2label.json'), 'w') as f:
        json.dump(id2label, f)
    print(f'Model saved to {args.model_dir}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--train_data', default='data/train.txt', help='训练数据路径')
    parser.add_argument('--label_list', default='data/label_list.txt', help='标签列表路径')
    parser.add_argument('--pretrained_model', default='bert-base-uncased', help='预训练模型名称或路径')
    parser.add_argument('--model_dir', default='models/bert_ner_model', help='模型保存路径')
    parser.add_argument('--epochs', type=int, default=3, help='训练轮数')
    parser.add_argument('--max_len', type=int, default=128, help='序列最大长度')
    parser.add_argument('--batch_size', type=int, default=16, help='批次大小')
    parser.add_argument('--lr', type=float, default=5e-5, help='学习率')
    args = parser.parse_args()
    train(args)

5.3 模型评估 (src/evaluate.py)

import argparse
import os
import json
import torch
from torch.utils.data import DataLoader
from transformers import BertForTokenClassification, BertTokenizer
from preprocess import NERDataset
from seqeval.metrics import classification_report

def load_labels(label_path):
    """
    加载标签列表，并创建标签与ID之间的映射。

    Args:
        label_path (str): 标签列表文件路径。

    Returns:
        labels (List[str]): 标签列表。
        label2id (dict): 标签到ID的映射。
        id2label (dict): ID到标签的映射。
    """
    with open(label_path, 'r') as f:
        labels = [line.strip() for line in f]
    label2id = {label: idx for idx, label in enumerate(labels)}
    id2label = {idx: label for idx, label in enumerate(labels)}
    return labels, label2id, id2label

def evaluate(args):
    """
    模型评估主函数。

    Args:
        args (argparse.Namespace): 命令行参数。
    """
    # 加载标签和分词器
    labels, label2id, id2label = load_labels(args.label_list)
    tokenizer = BertTokenizer.from_pretrained(args.model_dir)
    model = BertForTokenClassification.from_pretrained(args.model_dir)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    # 加载验证数据
    eval_dataset = NERDataset(args.eval_data, tokenizer, label2id, args.max_len)
    eval_loader = DataLoader(eval_dataset, batch_size=args.batch_size)

    # 模型评估
    all_preds, all_labels = [], []
    model.eval()
    with torch.no_grad():
        for batch in eval_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels']
            outputs = model(input_ids, attention_mask=attention_mask)
            logits = outputs.logits
            preds = torch.argmax(logits, dim=-1).cpu().numpy()
            labels = labels.numpy()
            for pred, label in zip(preds, labels):
                pred_labels = [id2label[p] for p, l in zip(pred, label) if l != -100]
                true_labels = [id2label[l] for p, l in zip(pred, label) if l != -100]
                all_preds.append(pred_labels)
                all_labels.append(true_labels)
    report = classification_report(all_labels, all_preds)
    print(report)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--eval_data', default='data/dev.txt', help='验证数据路径')
    parser.add_argument('--label_list', default='data/label_list.txt', help='标签列表路径')
    parser.add_argument('--model_dir', default='models/bert_ner_model', help='模型路径')
    parser.add_argument('--max_len', type=int, default=128, help='序列最大长度')
    parser.add_argument('--batch_size', type=int, default=16, help='批次大小')
    args = parser.parse_args()
    evaluate(args)

5.4 模型推理 (src/inference.py)

import argparse
import os
import json
import torch
from transformers import BertForTokenClassification, BertTokenizer

def load_labels(label_path):
    """
    加载标签列表，并创建ID到标签的映射。

    Args:
        label_path (str): 标签列表文件路径。

    Returns:
        id2label (dict): ID到标签的映射。
    """
    with open(label_path, 'r') as f:
        labels = [line.strip() for line in f]
    id2label = {idx: label for idx, label in enumerate(labels)}
    return id2label

def predict(args):
    """
    模型推理主函数。

    Args:
        args (argparse.Namespace): 命令行参数。
    """
    # 加载标签和分词器
    id2label = load_labels(os.path.join(args.model_dir, 'label_list.txt'))
    tokenizer = BertTokenizer.from_pretrained(args.model_dir)
    model = BertForTokenClassification.from_pretrained(args.model_dir)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()

    # 对输入文本进行分词和编码
    words = args.text.strip().split()
    encoding = tokenizer(
        words,
        is_split_into_words=True,
        return_offsets_mapping=True,
        padding='max_length',
        truncation=True,
        max_length=args.max_len,
        return_tensors='pt'
    )
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    # 模型推理
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=-1).cpu().numpy()[0]
    word_ids = encoding.word_ids()

    # 获取预测结果
    result = []
    for idx, word_id in enumerate(word_ids):
        if word_id is not None and word_id < len(words):
            result.append((words[word_id], id2label[predictions[idx]]))

    # 打印结果
    for word, label in result:
        print(f'{word}\t{label}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--text', required=True, help='输入文本')
    parser.add_argument('--model_dir', default='models/bert_ner_model', help='模型路径')
    parser.add_argument('--max_len', type=int, default=128, help='序列最大长度')
    args = parser.parse_args()
    predict(args)

6. 项目运行

6.1 一键运行脚本 (run.sh)

#!/bin/bash

# 训练模型
python src/train.py \
    --train_data data/train.txt \
    --label_list data/label_list.txt \
    --pretrained_model bert-base-uncased \
    --model_dir models/bert_ner_model \
    --epochs 3 \
    --max_len 128 \
    --batch_size 16 \
    --lr 5e-5

# 评估模型
python src/evaluate.py \
    --eval_data data/dev.txt \
    --label_list data/label_list.txt \
    --model_dir models/bert_ner_model \
    --max_len 128 \
    --batch_size 16

# 推理示例
python src/inference.py \
    --text "John lives in New York City." \
    --model_dir models/bert_ner_model \
    --max_len 128

注意：运行前请确保脚本具有执行权限。

chmod +x run.sh
./run.sh

6.2 手动运行

如果不使用一键脚本，可以手动执行以下命令。

6.2.1 训练模型
python src/train.py \
    --train_data data/train.txt \
    --label_list data/label_list.txt \
    --pretrained_model bert-base-uncased \
    --model_dir models/bert_ner_model \
    --epochs 3 \
    --max_len 128 \
    --batch_size 16 \
    --lr 5e-5
6.2.2 评估模型
python src/evaluate.py \
    --eval_data data/dev.txt \
    --label_list data/label_list.txt \
    --model_dir models/bert_ner_model \
    --max_len 128 \
    --batch_size 16
6.2.3 推理示例
python src/inference.py \
    --text "John lives in New York City." \
    --model_dir models/bert_ner_model \
    --max_len 128

7. 结果展示

7.1 训练日志

Epoch 1/3, Loss: 0.2453
Epoch 2/3, Loss: 0.1237
Epoch 3/3, Loss: 0.0784
Model saved to models/bert_ner_model

7.2 验证报告

              precision    recall  f1-score   support

       MISC       0.85      0.80      0.82        51
        PER       0.94      0.92      0.93        68
        ORG       0.89      0.86      0.87        59
        LOC       0.91      0.95      0.93        74

   micro avg       0.90      0.88      0.89       252
   macro avg       0.90      0.88      0.89       252
weighted avg       0.90      0.88      0.89       252

7.3 推理示例

输入文本：

John lives in New York City.

输出结果：

John    B-PER
lives   O
in      O
New     B-LOC
York    I-LOC
City.   I-LOC

8. 结论

本项目基于BERT模型，成功地实现了命名实体识别任务，完整展示了从数据预处理、模型训练、模型评估到模型推理的全过程。通过使用预训练语言模型，模型在NER任务中取得了较好的性能，证明了BERT在序列标注任务中的强大能力。

9. 参考资料

  • BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
  • Hugging Face Transformers Documentation
  • Seqeval: A Python framework for sequence labeling evaluation

http://www.kler.cn/a/396880.html

相关文章：

  • 二五、pxe自动装机
  • flutter字体大小切换案例 小字体，标准字体，大字体，超大字体案例
  • 解决背景图因为图片路径中携带括号导致图片无法显示的问题
  • git使用及上线流程（仅为我工作中常用）
  • tdengine学习笔记
  • java算法性能调优：详尽探讨时间复杂度与空间复杂度的分析与优化“
  • 网关的发展趋势，网关开源项目Higress
  • WebRTC 和 WebSocket
  • 选择排序+快速排序递归版（二）
  • 缓存冲突（Cache Conflict）
  • Spring：IOC实例化对象bean的方式
  • 蔚来Java面试题及参考答案
  • 冷启动 VS 热启动
  • 职场汇报技巧：选择合适的汇报形式与提供数据依据
  • C++ 的发展
  • ArkUI---使用弹窗---@ohos.promptAction (弹窗)
  • Linux 实现自动登陆远程机器
  • Qt之QTreeWidget通过撤销栈移除item
  • 软考之RESTful 架构的特点
  • uview Collapse折叠面板无法动态设置展开问题（微信小程序）
  • Docker在微服务架构中的应用
  • 算法之二分查找优化：leetcode34:在排序数组中查找元素的第一个和最后一个位置
  • 用 Python 从零开始创建神经网络（七）：梯度下降（Gradient Descent）/导数（Derivatives）
  • 27-压力测试
  • ASFSSA-VMD多策略改进的麻雀搜索算法优化变分模态分解
  • JavaWeb之AJAX