基于BERT的命名实体识别（NER）
目录
- 项目背景
- 项目结构
- 环境准备
- 数据准备
-
代码实现
- 5.1 数据预处理 (src/preprocess.py)
- 5.2 模型训练 (src/train.py)
- 5.3 模型评估 (src/evaluate.py)
- 5.4 模型推理 (src/inference.py)
-
项目运行
- 6.1 一键运行脚本 (run.sh)
- 6.2 手动运行
- 结果展示
- 结论
- 参考资料
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