命名实体识别概要
数据预处理
首先看数据结构，针对label提取所有的实体类型并且基于BMES做标签映射，将映射表保存。
python
import json
import os
def save_json(data, file_path):
"""
将数据以格式化的 JSON 形式保存到文件
"""
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=4)
def collect_entity_types_from_file(file_path):
"""
从单个数据文件中提取所有唯一的实体类型
"""
types = set()
with open(file_path, 'r', encoding='utf-8') as f:
all_data = json.load(f)
for data in all_data:
# 遍历实体列表，提取 'type' 字段
for entity in data['entities']:
types.add(entity['type'])
return types
def generate_tag_map(data_files, output_file):
"""
从数据文件构建 BMES 标签映射并保存
"""
# 1. 从所有文件中收集实体类型
all_entity_types = set()
for file_path in data_files:
all_entity_types.update(collect_entity_types_from_file(file_path))
# 2. 排序以保证映射一致性
sorted_types = sorted(list(all_entity_types))
print(f"发现的实体类型: {sorted_types}")
# 3. 构建 BMES 标签映射
tag_to_id = {'O': 0} # 'O' 代表非实体
for entity_type in sorted_types:
for prefix in ['B', 'M', 'E', 'S']:
tag_name = f"{prefix}-{entity_type}"
tag_to_id[tag_name] = len(tag_to_id)
print(f"\n已生成 {len(tag_to_id)} 个标签映射。")
# 4. 保存映射文件
save_json(tag_to_id, output_file)
print(f"标签映射已保存至: {output_file}")
if __name__ == '__main__':
# 定义输入的数据文件和期望的输出路径
train_file = './data/CMeEE-V2_train.json'
dev_file = './data/CMeEE-V2_dev.json'
output_path = './data/categories.json'
generate_tag_map(data_files=[train_file, dev_file], output_file=output_path)针对文本：规范化文本（统一全角半角），统计词频，删掉小于最小词频的单词，添加pad和unk，注意排序使词汇表可复现。
python
import json
import os
from collections import Counter
def save_json(data, file_path):
"""
将数据以易于阅读的格式保存为 JSON 文件
"""
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=4)
def normalize_text(text):
"""
规范化文本
"""
full_width = "０１２３４５６７８９ＡＢＣＤＥＦＧＨＩＪＫＬＭＮＯＰＱＲＳＴＵＶＷＸＹＺａｂｃｄｅｆｇｈｉｊｋｌｍｎｏｐｑｒｓｔｕｖｗｘｙｚ！＃＄％＆'（）＊＋，－．／：；＜＝＞？＠［＼］＾＿｀｛｜｝～＂"
half_width = r"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!#$%&'" + r'()*+,-./:;<=>?@[\]^_`{|}~".'
mapping = str.maketrans(full_width, half_width)
return text.translate(mapping)
def create_char_vocab(data_files, output_file, min_freq=1):
"""
从数据文件创建字符级词汇表
"""
char_counts = Counter()
for file_path in data_files:
with open(file_path, 'r', encoding='utf-8') as f:
all_data = json.load(f)
for data in all_data:
text = normalize_text(data['text'])
char_counts.update(list(text))
# 过滤低频词
frequent_chars = [char for char, count in char_counts.items() if count >= min_freq]
# 保证每次生成结果一致
frequent_chars.sort()
# 添加特殊标记
special_tokens = ["<PAD>", "<UNK>"]
final_vocab_list = special_tokens + frequent_chars
print(f"词汇表大小 (min_freq={min_freq}): {len(final_vocab_list)}")
# 保存词汇表
save_json(final_vocab_list, output_file)
print(f"词汇表已保存至: {output_file}")
if __name__ == '__main__':
train_file = './data/CMeEE-V2_train.json'
dev_file = './data/CMeEE-V2_dev.json'
output_path = './data/vocabulary.json'
# 设置字符最低频率，1表示包含所有出现过的字符
create_char_vocab(data_files=[train_file, dev_file], output_file=output_path, min_freq=1)构建vocabulary类，初始化生成token和id的映射，并且又__len__方法和token_id转化方法。
构建dataset，初始化传入文本，标签映射表，vocab类，实现__len__方法和__getitem__方法，其中__getitem__方法将输入的data转成token_id和标签id映射标识。返回token_ids和label_ids。
构建dataloader，首先构建collate_fn完成序列的pad和mask序列的构建，然后基于dataset生成dataloader。
python
import json
import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
def normalize_text(text):
"""
规范化文本，例如将全角字符转换为半角字符。
"""
full_width = "０１２３４５６７８９ＡＢＣＤＥＦＧＨＩＪＫＬＭＮＯＰＱＲＳＴＵＶＷＸＹＺａｂｃｄｅｆｇｈｉｊｋｌｍｎｏｐｑｒｓｔｕｖｗｘｙｚ！＃＄％＆'（）＊＋，－．／：；＜＝＞？＠［＼］＾＿｀｛｜｝～＂"
half_width = r"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!#$%&'" + r'()*+,-./:;<=>?@[\]^_`{|}~".'
mapping = str.maketrans(full_width, half_width)
return text.translate(mapping)
class Vocabulary:
"""
负责管理词汇表和 token 到 id 的映射。
"""
def __init__(self, vocab_path):
with open(vocab_path, 'r', encoding='utf-8') as f:
self.tokens = json.load(f)
self.token_to_id = {token: i for i, token in enumerate(self.tokens)}
self.pad_id = self.token_to_id['<PAD>']
self.unk_id = self.token_to_id['<UNK>']
def __len__(self):
return len(self.tokens)
def convert_tokens_to_ids(self, tokens):
return [self.token_to_id.get(token, self.unk_id) for token in tokens]
class NerDataset(Dataset):
"""
处理 NER 数据，并将其转换为适用于 PyTorch 模型的格式。
"""
def __init__(self, data_path, vocab: Vocabulary, tag_map: dict):
self.vocab = vocab
self.tag_to_id = tag_map
with open(data_path, 'r', encoding='utf-8') as f:
self.records = json.load(f)
def __len__(self):
return len(self.records)
def __getitem__(self, idx):
record = self.records[idx]
text = normalize_text(record['text'])
tokens = list(text)
# 将文本 tokens 转换为 ids
token_ids = self.vocab.convert_tokens_to_ids(tokens)
# 初始化标签序列为 'O'
tags = ['O'] * len(tokens)
for entity in record.get('entities', []):
entity_type = entity['type']
start = entity['start_idx']
end = entity['end_idx']
if end >= len(tokens): continue
if start == end:
tags[start] = f'S-{entity_type}'
else:
tags[start] = f'B-{entity_type}'
tags[end] = f'E-{entity_type}'
for i in range(start + 1, end):
tags[i] = f'M-{entity_type}'
# 将标签转换为 ids
label_ids = [self.tag_to_id.get(tag, self.tag_to_id['O']) for tag in tags]
return {
"token_ids": torch.tensor(token_ids, dtype=torch.long),
"label_ids": torch.tensor(label_ids, dtype=torch.long)
}
def create_ner_dataloader(data_path, vocab, tag_map, batch_size, shuffle=False):
"""
创建 NER 任务的 DataLoader。
"""
dataset = NerDataset(data_path, vocab, tag_map)
def collate_batch(batch):
token_ids_list = [item['token_ids'] for item in batch]
label_ids_list = [item['label_ids'] for item in batch]
padded_token_ids = pad_sequence(token_ids_list, batch_first=True, padding_value=vocab.pad_id)
padded_label_ids = pad_sequence(label_ids_list, batch_first=True, padding_value=-100) # -100 用于在计算损失时忽略填充部分
attention_mask = (padded_token_ids != vocab.pad_id).long()
return {
"token_ids": padded_token_ids,
"label_ids": padded_label_ids,
"attention_mask": attention_mask
}
return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_batch)
if __name__ == '__main__':
# 文件路径
train_file = './data/CMeEE-V2_train.json'
vocab_file = './data/vocabulary.json'
categories_file = './data/categories.json'
# 1. 加载词汇表和标签映射
vocabulary = Vocabulary(vocab_path=vocab_file)
with open(categories_file, 'r', encoding='utf-8') as f:
tag_map = json.load(f)
print("词汇表和标签映射加载完成。")
# 2. 创建 DataLoader
train_loader = create_ner_dataloader(
data_path=train_file,
vocab=vocabulary,
tag_map=tag_map,
batch_size=4,
shuffle=True
)
print("DataLoader 创建完成。")
# 3. 验证一个批次的数据
print("\n--- 验证一个批次的数据 ---")
batch = next(iter(train_loader))
print(f" Token IDs (shape): {batch['token_ids'].shape}")
print(f" Label IDs (shape): {batch['label_ids'].shape}")
print(f" Attention Mask (shape): {batch['attention_mask'].shape}")
print(f" Token IDs (sample): {batch['token_ids'][0][:20]}...")
print(f" Label IDs (sample): {batch['label_ids'][0][:20]}...")
print(f" Attention Mask (sample): {batch['attention_mask'][0][:20]}...")nn.ModuleList vs nn.Sequential
在 PyTorch 中，nn.ModuleList 和 nn.Sequential 都是用来容纳多个子模块的容器，但它们的设计思想和使用场景不同：
nn.Sequential：像一个自动化的流水线，数据会自动按顺序流过每一层。适用于简单的线性堆叠，但无法实现层间的复杂交互。- **
nn.ModuleList**：更像一个普通的 Python 列表 ，只负责存储模块，而不会自动执行它们。你需要在forward方法中手动编写循环来调用每一层，所以可以在层与层之间加入自定义逻辑（如残差连接）。
model构建
loss_fn = nn.CrossEntropyLoss(ignore_index=-100, reduction='none')
reduction='none'：正常损失返回的是一个值，设置该参数使损失返回每个标签上的损失，构成损失矩阵。
ignore_index 生效 ：可以看到 label_ids 中值为 -100 的填充位置，其对应的损失值为 0。如果用此方法，在单向模型中可以不传mask序列，但在双向模型中不行。
torch.nn.utils.rnn.pack_padded_sequence：这个函数的主要作用是接收一个 填充后 的 input 张量，以及一个记录了 真实长度 的 lengths 列表。它会返回一个 PackedSequence 对象，可以把它想象成一个"压缩"后的数据包，其中所有的填充位都被暂时移除了。RNN 模块在接收到这个特殊对象后，其内部就能正确、高效地处理变长序列。
当然，有"打包"就有"解包"。与之对应的 pad_packed_sequence 函数会负责将 RNN 计算完成后的 PackedSequence 对象再"解压"还原成带有填充的、规整的 Tensor。
python
import torch
import torch.nn as nn
import torch.nn.utils.rnn as rnn
class BiGRUNerNetWork(nn.Module):
def __init__(self, vocab_size, hidden_size, num_tags, num_gru_layers=1):
super().__init__()
# 1. Token Embedding 层
self.embedding = nn.Embedding(vocab_size, hidden_size)
# 2. 使用 ModuleList 构建多层双向 GRU
self.gru_layers = nn.ModuleList()
for _ in range(num_gru_layers):
self.gru_layers.append(
nn.GRU(
input_size=hidden_size,
hidden_size=hidden_size,
num_layers=1,
batch_first=True,
bidirectional=True # 开启双向
)
)
# 3. 特征融合层
self.fc = nn.Linear(hidden_size * 2, hidden_size)
# 4. 分类决策层 (Classifier)
self.classifier = nn.Linear(hidden_size, num_tags)
def forward(self, token_ids, attention_mask):
# 1. 计算真实长度
lengths = attention_mask.sum(dim=1).cpu()
# 2. 获取词向量
embedded_text = self.embedding(token_ids)
# 3. 打包序列
current_packed_input = rnn.pack_padded_sequence(
embedded_text, lengths, batch_first=True, enforce_sorted=False
)
# 4. 循环通过 GRU 层
for gru_layer in self.gru_layers:
# GRU 输出 (packed)
packed_output, _ = gru_layer(current_packed_input)
# 解包以进行后续操作，并指定 total_length
output, _ = rnn.pad_packed_sequence(
packed_output, batch_first=True, total_length=token_ids.shape[1]
)
# 特征融合
features = self.fc(output)
# 残差连接
# 同样需要解包上一层的输入
input_padded, _ = rnn.pad_packed_sequence(
current_packed_input, batch_first=True, total_length=token_ids.shape[1]
)
current_input = features + input_padded
# 重新打包作为下一层的输入
current_packed_input = rnn.pack_padded_sequence(
current_input, lengths, batch_first=True, enforce_sorted=False
)
# 5. 解包最终输出用于分类
final_output, _ = rnn.pad_packed_sequence(
current_packed_input, batch_first=True, total_length=token_ids.shape[1]
)
# 6. 分类
logits = self.classifier(final_output)
return logits
if __name__ == '__main__':
token_ids = torch.tensor([
[210, 18, 871, 147, 0, 0, 0, 0],
[922, 2962, 842, 210, 18, 871, 147, 0]
], dtype=torch.int64)
# attention_mask 标记哪些是真实 token (1) 哪些是填充 (0)
attention_mask = torch.tensor([
[1, 1, 1, 1, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 1, 1, 0]
], dtype=torch.int64)
label_ids = torch.tensor([
[0, 0, 0, 0, -100, -100, -100, -100],
[0, 0, 0, 0, 0, 0, 0, -100]
], dtype=torch.int64)
# 实例化模型
model = BiGRUNerNetWork(
vocab_size=10000,
hidden_size=128,
num_tags=37,
num_gru_layers=2
)
# 3. 执行前向传播
logits = model(token_ids=token_ids, attention_mask=attention_mask)
# 4. 构造损失函数
loss_fn = nn.CrossEntropyLoss(ignore_index=-100, reduction='none')
# 5. 计算损失
# CrossEntropyLoss 要求类别维度在前，所以需要交换最后两个维度
# [batch, seq_len, num_tags] -> [batch, num_tags, seq_len]
permuted_logits = torch.permute(logits, dims=(0, 2, 1))
loss = loss_fn(permuted_logits, label_ids)
# 6. 打印结果
print(f"Logits shape: {logits.shape}")
print(f"Loss shape: {loss.shape}")
print("\n每个 Token 的损失:")
print(loss)组件构建
Trainer
在开始编写 Trainer 类之前，先在 src/ 目录下创建一个 trainer 文件夹，并在其中新建一个 trainer.py 文件，用于存放 Trainer 类的定义。
Trainer 只负责"训练" : Trainer 类的核心职责是执行标准的训练和评估循环。
init：初始化model: PyTorch 模型。 optimizer: 优化器。 loss_fn: 损失函数。 train_loader: 训练数据加载器。 dev_loader: 验证数据加载器。 eval_metric_fn: 评估函数。 output_dir: 模型输出目录。 device: 训练设备。
fit(self, epochs):训练的主入口
_train_one_epoch(self):封装一个epoch训练流程
_train_step(self, batch):封装一个训练步骤的逻辑（前向、损失、反向）
_evaluate(self):封装评估逻辑。
_evaluation_step(self, batch):封装一个评估步骤的逻辑（前向、损失）。
_save_checkpoint(self, is_best=False):封装模型保存逻辑。
在我们当前的"组件式组装"设计中，虽然 Trainer 不直接接收整个 config 对象（以保持解耦），但 config 依然是所有"零件"的参数来源。
python
import torch
from tqdm import tqdm
import os
from dataclasses import asdict
class Trainer:
def __init__(self, model, optimizer, loss_fn, train_loader, dev_loader=None,
eval_metric_fn=None, output_dir=None, device='cpu'):
self.model = model.to(device)
self.optimizer = optimizer
self.loss_fn = loss_fn
self.train_loader = train_loader
self.dev_loader = dev_loader
self.eval_metric_fn = eval_metric_fn
self.output_dir = output_dir
self.device = torch.device(device)
if self.output_dir:
os.makedirs(self.output_dir, exist_ok=True)
def fit(self, epochs):
best_metric = float('inf') # 初始化一个无穷大的 best_metric，用于后续比较
for epoch in range(1, epochs + 1):
# 1. 执行一个周期的训练
train_loss = self._train_one_epoch()
print(f"Epoch {epoch} - Training Loss: {train_loss:.4f}")
# 2. 执行评估
metrics = self._evaluate()
if metrics:
print(f"Epoch {epoch} - Validation Metrics: {metrics}")
current_metric = metrics.get('loss') # 默认监控验证集 loss
# 3. 如果当前 metric 优于历史最优，则保存最佳模型
if current_metric < best_metric:
best_metric = current_metric
if self.output_dir:
self._save_checkpoint(is_best=True)
print(f"New best model saved with validation loss: {best_metric:.4f}")
# 4. 每个 epoch 结束后，保存最新的模型状态
if self.output_dir:
self._save_checkpoint(is_best=False)
def _train_one_epoch(self):
"""执行一个完整的训练周期。"""
self.model.train() # 设置为训练模式
total_loss = 0
# 使用 tqdm 显示进度条
for batch in tqdm(self.train_loader, desc=f"Training Epoch"):
outputs = self._train_step(batch)
total_loss += outputs['loss'].item() # 累加 loss
return total_loss / len(self.train_loader) # 返回平均 loss
def _train_step(self, batch):
"""执行单个训练步骤（前向、损失、反向）。"""
# 1. 将数据移动到指定设备
batch = {k: v.to(self.device) for k, v in batch.items() if isinstance(v, torch.Tensor)}
# 2. 模型前向传播
logits = self.model(token_ids=batch['token_ids'], attention_mask=batch['attention_mask'])
# 3. 计算损失
# CrossEntropyLoss 要求 logits 的形状为 [B, C, L]，label_ids 的形状为 [B, L]
loss = self.loss_fn(logits.permute(0, 2, 1), batch['label_ids'])
# 4. 反向传播与参数更新
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {'loss': loss, 'logits': logits}
def _evaluate(self):
"""在验证集上执行评估。"""
if self.dev_loader is None:
return None
self.model.eval() # 设置为评估模式
total_loss = 0
all_logits = []
all_labels = []
all_attention_mask = []
with torch.no_grad(): # 禁用梯度计算
for batch in tqdm(self.dev_loader, desc="Evaluating"):
outputs = self._evaluation_step(batch)
total_loss += outputs['loss'].item()
# 收集所有批次的 logits 和 labels，用于后续评估
all_logits.append(outputs['logits'].cpu())
all_labels.append(batch['label_ids'].cpu())
all_attention_mask.append(batch['attention_mask'].cpu())
metrics = {}
# 如果提供了评估函数，则调用它来计算指标
if self.eval_metric_fn:
metrics = self.eval_metric_fn(all_logits, all_labels, all_attention_mask)
# 计算并记录平均 loss
metrics['loss'] = total_loss / len(self.dev_loader)
return metrics
def _evaluation_step(self, batch):
"""执行单个评估步骤（前向、损失）。"""
# 1. 将数据移动到指定设备
batch = {k: v.to(self.device) for k, v in batch.items() if isinstance(v, torch.Tensor)}
# 2. 模型前向传播
logits = self.model(token_ids=batch['token_ids'], attention_mask=batch['attention_mask'])
# 3. 计算损失
loss = self.loss_fn(logits.permute(0, 2, 1), batch['label_ids'])
return {'loss': loss, 'logits': logits}
def _save_checkpoint(self, is_best):
"""保存模型检查点。"""
state = {'model_state_dict': self.model.state_dict()}
if is_best:
# 保存最佳模型
torch.save(state, os.path.join(self.output_dir, 'best_model.pth'))
# 保存最新模型
torch.save(state, os.path.join(self.output_dir, 'last_model.pth'))配置类config参数
python
# src/configs/configs.py
import torch
from dataclasses import dataclass, field
@dataclass
class NerConfig:
# --- 路径参数 ---
data_dir: str = "data"
train_file: str = "CMeEE-V2_train.json"
dev_file: str = "CMeEE-V2_dev.json"
vocab_file: str = "vocabulary.json"
tags_file: str = "categories.json"
output_dir: str = "output"
# --- 训练参数 ---
batch_size: int = 32
epochs: int = 20
learning_rate: float = 1e-3
device: str = field(default_factory=lambda: 'cuda' if torch.cuda.is_available() else 'cpu')
# --- 模型参数 ---
hidden_size: int = 256
num_gru_layers: int = 2@dataclass 是 Python 3.7 引入的装饰器，可以简化类的编写。对于 TrainerConfig 这样的配置类，它会自动生成构造函数 (__init__)，无需再手动编写冗长的参数赋值代码。同时，它还会生成一个友好的打印格式 (__repr__)，这意味着 print(config) 会清晰地展示所有参数和值，便于调试。
模型组件
第一步：创建模型目录
在 src/ 目录下创建一个新的文件夹 models。
第二步：定义模型基类
在构建具体的模型之前，可以先在 src/models/ 目录下创建一个 base.py 文件来定义一个 模型基类 。这个基类使用 Python 的 abc 模块（Abstract Base Classes）来规定所有 NER 模型都必须遵循的一个统一接口。
这样做的好处是：
- 强制接口统一 ：所有模型都必须实现一个
forward方法，且接收相同的参数（token_ids,attention_mask）。这保证了Trainer可以与任何我们未来创建的新模型（如 BERT-NER, LSTM-NER）无缝协作，无需修改Trainer的代码。 - 提高可读性与可维护性：代码的结构更清晰，别人接手项目时，只需查看基类就能明白模型部分的接口规范。
python
# src/models/base.py
import torch.nn as nn
from abc import ABC, abstractmethod
class BaseNerNetwork(nn.Module, ABC):
@abstractmethod
def forward(self, token_ids, attention_mask):
"""
定义所有 NER 模型都必须遵循的前向传播接口。
Args:
token_ids (torch.Tensor): [batch_size, seq_len]
attention_mask (torch.Tensor): [batch_size, seq_len]
Returns:
torch.Tensor: Logits, [batch_size, seq_len, num_tags]
"""
raise NotImplementedError
python
# src/models/ner_model.py
import torch.nn as nn
import torch.nn.utils.rnn as rnn
from .base import BaseNerNetwork # 导入基类
class BiGRUNerNetWork(BaseNerNetwork): # 继承自 BaseNerNetwork
# ... (省略具体实现，与前文一致) ...数据加载组件
定义 NerDataset 类，DataLoader
分词器组件
构建分词器基类，在换其他分词器时保障后续接口可以使用。
python
# src/tokenizer/base.py
from abc import ABC, abstractmethod
class BaseTokenizer(ABC):
@abstractmethod
def text_to_tokens(self, text: str) -> list[str]:
"""将文本分割成 token 列表。"""
raise NotImplementedError
@abstractmethod
def tokens_to_ids(self, tokens: list[str]) -> list[int]:
"""将 token 列表转换为 ID 列表。"""
raise NotImplementedError
def encode(self, text: str) -> list[int]:
"""将文本直接编码为 ID 列表的便捷方法。"""
tokens = self.text_to_tokens(text)
return self.tokens_to_ids(tokens)
@abstractmethod
def get_pad_id(self) -> int:
"""获取填充 token 的 ID。"""
raise NotImplementedError
python
# src/tokenizer/char_tokenizer.py
from .vocabulary import Vocabulary
from .base import BaseTokenizer
def normalize_text(text):
# ... (省略 normalize_text 函数实现) ...
class CharTokenizer(BaseTokenizer):
def __init__(self, vocab: Vocabulary):
self.vocab = vocab
def text_to_tokens(self, text: str):
normalized_text = normalize_text(text)
return list(normalized_text)
def tokens_to_ids(self, tokens: list[str]):
return self.vocab.convert_tokens_to_ids(tokens)
def get_pad_id(self) -> int:
return self.vocab.pad_id评估组件
def _trans_entity2tuple(label_ids, id2tag):将标签ID序列转换为实体元组列表（严格 BMES 解码）。仅在遇到 E- 或 S- 时落盘；遇到新的 B- 或 O 不闭合未完成片段。
当前 calculate_entity_level_metrics 的实现，在面对 Trainer 的批量评估场景时，会遇到两个问题：
-
处理填充 ：在一个批次中，不同长度的句子会被填充到相同长度。这些填充位（Padding）不应参与评估。我们需要利用
attention_mask机制，来过滤掉所有因填充而产生的无效 Token，确保评估只在有效的序列片段上进行。 -
追踪样本来源 ：当处理一个批次的多个样本时，必须能区分每个实体到底来自哪个样本。例如，批次中的第一个样本和第二个样本可能在相同的位置
(0, 2)都有一个'dis'类型的实体。如果在解码时不加以区分，这两个独立的实体在存入set时会被误判为同一个。为了准确区分来自同一批次中不同样本的实体，设计了一种方案：为每个解码出的实体附加其所在样本的唯一ID（即批次内索引i）。确保每个实体都由一个唯一的(样本ID, 实体类型, 起始位置, 结束位置)四元组来标识，从根本上解决实体归属混淆的问题。
python
# src/metrics/entity_metrics.py
import torch
def _trans_entity2tuple(label_ids, id2tag):
"""
将标签ID序列转换为实体元组列表（严格 BMES 解码）。
仅在遇到 E- 或 S- 时落盘；遇到新的 B- 或 O 不闭合未完成片段。
"""
entities = []
current_entity = None
for i, label_id in enumerate(label_ids):
# 将标签ID映射为字符串标签，未知则视作 'O'
tag = id2tag.get(label_id.item(), 'O')
if tag.startswith('B-'):
# 开启新片段：记录类型与起始位置；end 暂定为 i+1
current_entity = (tag[2:], i, i + 1)
elif tag.startswith('M-'):
# 仅当已存在片段，且类型一致时续接（扩展 end）
if current_entity and current_entity[0] == tag[2:]:
current_entity = (current_entity[0], current_entity[1], i + 1)
else:
# 类型不一致或不存在片段：丢弃未完成片段
current_entity = None
elif tag.startswith('E-'):
# 仅当已存在片段且类型一致时闭合并落盘
if current_entity and current_entity[0] == tag[2:]:
current_entity = (current_entity[0], current_entity[1], i + 1)
entities.append(current_entity)
# 无论是否匹配，E- 都视为一次片段结束
current_entity = None
elif tag.startswith('S-'):
# 单字实体：直接落盘（start=i, end=i+1）
entities.append((tag[2:], i, i + 1))
current_entity = None
else: # 'O'
# 非实体位置：严格模式不闭合未完成片段，直接丢弃
current_entity = None
# 返回集合去重
return set(entities)
def calculate_entity_level_metrics(all_pred_ids, all_label_ids, all_masks, id2tag):
"""
计算实体级别的精确率、召回率和 F1 分数。
"""
true_entities = set()
pred_entities = set()
sample_idx = 0
# 按批次遍历，同时保持 preds/labels/masks 对齐
for preds_batch, labels_batch, masks_batch in zip(all_pred_ids, all_label_ids, all_masks):
B = labels_batch.shape[0] # 当前批次样本数
for b in range(B):
# 对单个样本应用布尔掩码，去除 padding 位置
row_mask = masks_batch[b].bool()
row_labels = labels_batch[b][row_mask]
row_preds = preds_batch[b][row_mask]
# 严格 BMES 解码为实体集合
te = _trans_entity2tuple(row_labels, id2tag)
pe = _trans_entity2tuple(row_preds, id2tag)
# 为每个实体附加 (sample_idx,) 前缀，确保不同样本的相同实体不冲突
true_entities.update({(sample_idx,) + e for e in te})
pred_entities.update({(sample_idx,) + e for e in pe})
sample_idx += 1
num_correct = len(true_entities.intersection(pred_entities))
num_true = len(true_entities)
num_pred = len(pred_entities)
precision = num_correct / num_pred if num_pred > 0 else 0.0
recall = num_correct / num_true if num_true > 0 else 0.0
f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0
return {"precision": precision, "recall": recall, "f1": f1}模型推理与优化
解码预测序列：
模型的前向传播最终输出的是一个 logits 张量，形状为 [batch_size, seq_len, num_tags]。经过 argmax 操作后，会得到一个标签 ID 序列，例如 [0, 9, 10, 11, 0, ...]。
这个序列本身并不直观。为了进行实体级评估，或者将预测结果呈现给用户，必须实现一个 解码 (Decode) 函数，将这个数字序列转换成一个包含具体实体信息的列表，例如：[{"text": "高血压", "type": "dis", "start": 3, "end": 6}]。这个解码过程的核心，就是根据 BMES 标注体系的规则，从标签序列中解析出实体的边界和类型。
解码策略：
当前采用的是一种 "严格"模式 。任何不符合规范的序列（例如只有 B- 没有 E- 的实体）都会被直接放弃。这是最常见的做法，因为它能保证输出实体的规范性。
在某些特定的业务场景下，也可以采用更 "宽松"的策略 。例如，如果模型预测出一个 B-M-O 的序列，可以选择将 B-M 这部分作为一个实体输出，而不是完全丢弃它。这种策略的选择，取决于具体应用对"召回率"和"精确率"的不同侧重，需要根据实际需求来决定。
NerPredictor 核心流程
__init__ 方法的目标是加载并准备好所有推理所需的组件。
加载配置 : 从模型目录加载 config.json，获取模型超参数和相关文件路径
加载词汇表和标签映射 : 根据配置文件中的路径，加载 vocabulary.json 和 tags.json，并构建 id2tag 映射。
加载分词器 : 初始化 CharTokenizer。
初始化模型并加载权重:
- 根据配置实例化
BiGRUNerNetWork模型。 - 从模型目录加载
best_model.pth模型权重。这里需要使用map_location=self.device来确保模型可以被加载到指定的设备上（无论是 CPU 还是 GPU）。 - 调用
model.to(self.device)将模型移至指定设备。 - 调用
model.eval()将模型切换到评估模式，关闭 Dropout 和 BatchNorm 等只在训练时使用的层，确保预测结果的确定性。
python
# code/C8/06_predict.py
import torch
import json
import os
import argparse
from src.models.ner_model import BiGRUNerNetWork
from src.tokenizer.vocabulary import Vocabulary
from src.tokenizer.char_tokenizer import CharTokenizer
from src.utils.file_io import load_json
class NerPredictor:
def __init__(self, model_dir, device='cpu'):
self.device = torch.device(device)
# --- 1. 加载配置文件以获取模型参数 ---
config_path = os.path.join(model_dir, 'config.json')
self.config = load_json(config_path)
# --- 2. 加载词汇表和标签映射 ---
vocab_path = os.path.join(self.config["data_dir"], self.config["vocab_file"])
tags_path = os.path.join(self.config["data_dir"], self.config["tags_file"])
self.vocab = Vocabulary.load_from_file(vocab_path)
self.tokenizer = CharTokenizer(self.vocab)
tag_map = load_json(tags_path)
self.id2tag = {v: k for k, v in tag_map.items()}
# --- 3. 初始化模型并加载权重 ---
self.model = BiGRUNerNetWork(
vocab_size=len(self.vocab),
hidden_size=self.config["hidden_size"],
num_tags=len(tag_map),
num_gru_layers=self.config["num_gru_layers"]
)
model_path = os.path.join(model_dir, 'best_model.pth')
self.model.load_state_dict(torch.load(model_path, map_location=self.device)['model_state_dict'])
self.model.to(self.device)
self.model.eval()
def predict(self, text):
tokens = self.tokenizer.text_to_tokens(text)
token_ids = self.tokenizer.tokens_to_ids(tokens)
# --- 预处理 ---
token_ids_tensor = torch.tensor([token_ids], dtype=torch.long).to(self.device)
attention_mask = torch.ones_like(token_ids_tensor)
# --- 模型预测 ---
with torch.no_grad():
logits = self.model(token_ids_tensor, attention_mask)
# --- 后处理 ---
predictions = torch.argmax(logits, dim=-1).squeeze(0)
tags = [self.id2tag[id_.item()] for id_ in predictions]
return self._extract_entities(tokens, tags)
def _extract_entities(self, tokens, tags):
entities = []
current_entity = None
for i, tag in enumerate(tags):
if tag.startswith('B-'):
if current_entity:
pass
current_entity = {"text": tokens[i], "type": tag[2:], "start": i}
elif tag.startswith('M-'):
if current_entity and current_entity["type"] == tag[2:]:
current_entity["text"] += tokens[i]
else:
current_entity = None
elif tag.startswith('E-'):
if current_entity and current_entity["type"] == tag[2:]:
current_entity["text"] += tokens[i]
current_entity["end"] = i + 1
entities.append(current_entity)
current_entity = None
elif tag.startswith('S-'):
current_entity = None
entities.append({"text": tokens[i], "type": tag[2:], "start": i, "end": i + 1})
else: # 'O' 标签
current_entity = None
return entities
def main():
parser = argparse.ArgumentParser(description="NER Prediction")
parser.add_argument("--model_dir", type=str, required=True, help="Directory of the saved model and config.")
parser.add_argument("--text", type=str, required=True, help="Text to predict.")
args = parser.parse_args()
predictor = NerPredictor(model_dir=args.model_dir)
entities = predictor.predict(args.text)
print(f"Text: {args.text}")
print(f"Entities: {json.dumps(entities, ensure_ascii=False, indent=2)}")
if __name__ == "__main__":
main()自定义损失函数
加权交叉熵损失
最简单的方法就是"加权"。给数量稀少的实体标签（B, M, E, S）一个更高的权重，给数量庞大的非实体标签（O）一个较低的权重。例如，我们可以设置实体损失的权重为 10，非实体损失的权重为 1。这样，模型在反向传播时，如果弄错了一个实体 Token，会受到比弄错一个非实体 Token 大 10 倍的"惩罚"，从而迫使模型更加关注对实体的识别。
硬负样本挖掘
另一种思路是"采样"。在大量的非实体样本中，大部分是模型可以轻易正确预测的"简单样本"，它们对损失的贡献很小，反复学习意义不大。真正有价值的是那些模型容易搞错的"硬负样本"，例如一个模型倾向于预测为实体的非实体 Token。
硬负样本挖掘的做法是：在计算非实体部分的损失时，不计算所有非实体 Token 的平均损失，而是只选择其中损失值最大（Top-K）的一部分进行计算和反向传播。这样就相当于从海量的"多数派"中，筛选出了最有价值的"疑难样本"进行学习，提升了训练的效率和效果。
python
# code/C8/src/loss/ner_loss.py
import torch
import torch.nn as nn
class NerLoss(nn.Module):
"""
自定义 NER 损失函数，集成两种策略来对抗数据不均衡问题：
1. 加权交叉熵
2. 硬负样本挖掘
"""
def __init__(self, loss_type='cross_entropy', entity_weight=10.0, hard_negative_ratio=0.5, ignore_index=-100):
super().__init__()
# --- 参数定义 ---
self.loss_type = loss_type # 损失类型: 'cross_entropy', 'weighted_ce', 'hard_negative_mining'
self.entity_weight = entity_weight # 实体损失的权重
self.hard_negative_ratio = hard_negative_ratio # 硬负样本与正样本的比例
# 基础损失函数，设置为 'none' 模式以获取每个 token 的单独损失
self.base_loss_fn = nn.CrossEntropyLoss(reduction='none', ignore_index=ignore_index)
def forward(self, logits, labels):
"""
根据初始化时选择的 loss_type 计算损失。
"""
if self.loss_type == 'weighted_ce':
return self._weighted_cross_entropy(logits, labels)
elif self.loss_type == 'hard_negative_mining':
return self._hard_negative_mining(logits, labels)
else:
# 默认使用 PyTorch 原生的交叉熵损失
return self.base_loss_fn(logits, labels).mean()
def _weighted_cross_entropy(self, logits, labels):
"""
加权交叉熵损失的实现。
"""
# 计算每个 token 的基础损失, shape: [batch_size, seq_len]
loss_per_token = self.base_loss_fn(logits, labels)
# 创建掩码来区分实体和非实体 token
entity_mask = (labels > 0).float() # 实体 (B, M, E, S)
non_entity_mask = (labels == 0).float() # 非实体 (O)
# 分别计算实体和非实体部分的平均损失
entity_loss = torch.sum(loss_per_token * entity_mask) / (torch.sum(entity_mask) + 1e-8)
non_entity_loss = torch.sum(loss_per_token * non_entity_mask) / (torch.sum(non_entity_mask) + 1e-8)
# 根据预设权重，组合两部分损失
total_loss = self.entity_weight * entity_loss + 1.0 * non_entity_loss
return total_loss, entity_loss.detach(), non_entity_loss.detach()
def _hard_negative_mining(self, logits, labels):
"""
硬负样本挖掘损失的实现。
"""
# 计算每个 token 的基础损失
loss_per_token = self.base_loss_fn(logits, labels)
# 实体部分的损失计算与加权交叉熵方法相同
entity_mask = (labels > 0).float()
entity_loss = torch.sum(loss_per_token * entity_mask) / (torch.sum(entity_mask) + 1e-8)
# 筛选出所有非实体 token 的损失
non_entity_mask = (labels == 0).float()
non_entity_loss = loss_per_token * non_entity_mask
# 确定要挖掘的硬负样本数量
num_entities = torch.sum(entity_mask).item()
num_hard_negatives = int(num_entities * self.hard_negative_ratio)
# 如果当前批次没有实体，则按固定比例选择负样本，避免数量为0
if num_hard_negatives == 0:
num_non_entities = torch.sum(non_entity_mask).item()
num_hard_negatives = int(num_non_entities * 0.1)
# 从非实体损失中选出最大的 top-k 个作为硬负样本
topk_losses, _ = torch.topk(non_entity_loss.view(-1), k=num_hard_negatives)
# 计算硬负样本的平均损失
hard_negative_loss = torch.mean(topk_losses)
# 结合实体损失和硬负样本损失
total_loss = self.entity_weight * entity_loss + 1.0 * hard_negative_loss
return total_loss, entity_loss.detach(), hard_negative_loss.detach()