命名实体识别项目

这是一个基于 PyTorch 的命名实体识别（Named Entity Recognition, NER）项目，支持使用 BiLSTM 或 Transformer Encoder 进行序列标注训练，并通过 K-Fold 交叉验证评估模型表现。

项目当前提供：

按行标注数据读取与词表/标签表构建
固定长度序列 padding / truncation
BiLSTM 与 Transformer 两种 NER 模型
K-Fold 训练流程
验证集 accuracy、测试集分类报告与混淆矩阵可视化
基础单元测试

项目结构

Named-Entity-Recognition/
├── cli.py
├── ner_dataset.txt
├── requirements.txt
├── README.md
├── src/
│   ├── __init__.py
│   └── ner/
│       ├── __init__.py
│       ├── data.py
│       ├── data_utils.py
│       ├── evaluate.py
│       ├── models.py
│       └── train.py
└── tests/
    ├── sample_dataset.txt
    └── test_data.py

模块说明

文件	说明
`cli.py`	命令行入口，串联数据读取、建表、K-Fold 训练和最终评估
`src/ner/data.py`	主训练流程使用的数据模块，包含 `NERDataset` 和 K-Fold DataLoader 构建
`src/ner/data_utils.py`	测试中使用的轻量数据工具函数
`src/ner/models.py`	`BiLSTMTagger`、`TransformerTagger` 和位置编码实现
`src/ner/train.py`	训练循环、验证集评估、loss/accuracy 曲线展示
`src/ner/evaluate.py`	测试评估，输出实体标签分类报告并展示混淆矩阵
`tests/test_data.py`	数据读取与词表构建的基础测试

环境安装

建议使用虚拟环境：

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

依赖包括：

torch
numpy
scikit-learn
matplotlib
seaborn
pytest

数据格式

数据文件采用常见的 token-level 标注格式：

每行至少包含一个 token 和一个标签
代码读取每行的第一个字段作为 token，最后一个字段作为标签
句子之间使用空行分隔
标签通常使用 IOB 格式，例如 B-PER、I-PER、B-ORG、O

示例：

John B-PER
Doe I-PER
works O
at O
Google B-ORG

Jane B-PER
lives O
in O
New B-LOC
York I-LOC

项目根目录中的 ner_dataset.txt 可作为训练数据示例，tests/sample_dataset.txt 可作为测试数据示例。

训练模型

使用 BiLSTM：

python cli.py --data ner_dataset.txt --model bilstm --epochs 10 --folds 5

使用 Transformer：

python cli.py --data ner_dataset.txt --model transformer --epochs 10 --folds 5

命令行参数：

参数	必需	默认值	说明
`--data`	是	无	数据集路径
`--model`	否	`bilstm`	模型类型，可选 `bilstm` 或 `transformer`
`--epochs`	否	`3`	每一折训练 epoch 数
`--folds`	否	`5`	K-Fold 交叉验证折数

训练流程会先构建词表和标签表，再按 --folds 生成训练集/验证集 DataLoader。每一折都会重新初始化模型并训练，最终使用最后一折的验证集作为测试集调用 evaluate_on_test。

当前实现细节

数据处理

词表固定包含 <PAD> 和 <UNK>，索引分别为 0 和 1
标签表按标签首次出现顺序分配索引
NERDataset 默认 max_len=100
序列短于 max_len 时会 padding，长于 max_len 时会截断
padding token 使用 <PAD> 的索引 0
padding 标签使用索引 0，训练损失通过 ignore_index=tag_pad_idx 忽略该位置
K-Fold 使用 shuffle=True 和 random_state=42
DataLoader 默认 batch_size=32

BiLSTMTagger

默认结构：

nn.Embedding(vocab_size, 128, padding_idx=0)
单层双向 LSTM，总 hidden size 为 256
Dropout(0.3)
Linear(256, tagset_size)

输入 shape 为 (batch_size, seq_len)，输出 shape 为 (batch_size, seq_len, tagset_size)。

TransformerTagger

默认结构：

nn.Embedding(vocab_size, 128, padding_idx=0)
正弦/余弦位置编码
nn.TransformerEncoderLayer(d_model=128, nhead=8)
2 层 nn.TransformerEncoder
Linear(128, tagset_size)

输入 shape 为 (batch_size, seq_len)，内部会转换为 PyTorch Transformer 默认的 (seq_len, batch_size, embedding_dim) 格式，输出再转换回 (batch_size, seq_len, tagset_size)。

评估与可视化

训练时每个 epoch 会输出：

训练 loss
训练 accuracy
验证 accuracy

训练结束后，train.py 会通过 matplotlib.pyplot.show() 展示：

loss 曲线
accuracy 曲线

最终评估时，evaluate.py 会：

对非 padding 的 token 进行预测收集
输出实体标签的 classification_report，默认排除 O 标签
使用 Seaborn 展示包含所有标签的混淆矩阵

当前实现会展示图表，但不会自动保存图片文件。

运行测试

pytest tests/test_data.py -v

也可以直接运行全部测试：

pytest

Python 中使用

from src.ner.data import read_data, build_vocab, build_tag_vocab
from src.ner.models import BiLSTMTagger

sentences, tags = read_data("ner_dataset.txt")
word2idx = build_vocab(sentences)
tag2idx = build_tag_vocab(tags)

model = BiLSTMTagger(
    vocab_size=len(word2idx),
    tagset_size=len(tag2idx),
)

完整训练流程可参考 cli.py。

注意事项

当前 CLI 没有保存模型 checkpoint，需要保存模型时可在训练后自行调用 torch.save
当前评估使用最后一折验证集作为最终测试集，并没有单独的 train/validation/test 三份划分
当前 Transformer 没有显式传入 padding mask，因此 padding token 仍可能参与注意力计算
lengths 在数据集中记录的是原始长度；当句子超过 max_len 被截断时，当前代码仍保留原始长度
如果在无图形界面的环境运行，plt.show() 可能需要额外配置 Matplotlib backend

许可证

MIT License

作者

Jamie0807

FilesExpand file tree

README.md

Latest commit

History