跟着AI学AI - 诊断结论信息抽取 - 模型训练

数据增强和自动标注只是准备数据的手段，最终目的是用这些数据训练出好用的模型。

seqeval 是一个专门用于评估序列标注任务（如命名实体识别 NER）的 Python 库。它支持多种标签格式（如 BIO、IOBES 等），并提供了计算准确率（Accuracy）、查准率（Precision）、召回率（Recall）和 F1 分数等指标的便捷方法。seqeval 适用于命名实体识别、词性标注和语义角色标注等任务。

# 设置镜像源的环境变量
(vippython) PS D:\OpenSource\Python\VipPython> $env:HF_ENDPOINT = "https://hf-mirror.com"
# 安装依赖
(vippython) PS D:\OpenSource\Python\VipPython> uv pip install seqeval
(vippython) PS D:\OpenSource\Python\VipPython> uv pip install accelerate>=0.26.0
# 切换下目录，否则会报文件不存在
(vippython) PS D:\OpenSource\Python\VipPython> cd D:\OpenSource\Python\VipPython\information_extraction
(vippython) PS D:\OpenSource\Python\VipPython\information_extraction> uv run .\information_extraction\train_ner_model.py

train_ner_model.py

# train_ner_model.py
import json
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification
)
from torch.utils.data import Dataset
import numpy as np
from sklearn.metrics import classification_report, confusion_matrix
import os


class NERDataset(Dataset):
    def __init__(self, data, label2id):
        self.data = data
        self.label2id = label2id

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        return {
            'input_ids': torch.tensor(item['input_ids']),
            'attention_mask': torch.tensor(item['attention_mask']),
            'labels': torch.tensor(item['labels'])
        }


def compute_metrics(pred):
    """计算评估指标"""
    predictions, labels = pred
    predictions = np.argmax(predictions, axis=2)

    # 移除忽略的索引
    true_predictions = [
        [p for (p, l) in zip(pred_row, label_row) if l != -100]
        for pred_row, label_row in zip(predictions, labels)
    ]
    true_labels = [
        [l for l in label_row if l != -100]
        for label_row in labels
    ]

    # 计算指标
    from seqeval.metrics import classification_report as seq_report
    from seqeval.metrics import accuracy_score, f1_score, precision_score, recall_score

    # 将ID转换为标签名
    id2label = {v: k for k, v in label2id.items()}

    predictions_labels = [[id2label[p] for p in pred] for pred in true_predictions]
    references_labels = [[id2label[l] for l in ref] for ref in true_labels]

    return {
        'f1': f1_score(references_labels, predictions_labels),
        'precision': precision_score(references_labels, predictions_labels),
        'recall': recall_score(references_labels, predictions_labels),
        'accuracy': accuracy_score(references_labels, predictions_labels)
    }


def train_model(train_file, valid_file, label2id, output_dir='./ner_model'):
    """训练NER模型"""
    print("=" * 60)
    print("开始训练NER模型")
    print("=" * 60)

    # 1. 加载数据
    print("加载训练数据...")
    with open(train_file, 'r', encoding='utf-8') as f:
        train_data = json.load(f)

    with open(valid_file, 'r', encoding='utf-8') as f:
        valid_data = json.load(f)

    print(f"训练集: {len(train_data)} 条")
    print(f"验证集: {len(valid_data)} 条")

    # 2. 创建数据集
    train_dataset = NERDataset(train_data, label2id)
    valid_dataset = NERDataset(valid_data, label2id)

    # 3. 加载模型
    tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
    model = AutoModelForTokenClassification.from_pretrained(
        "bert-base-chinese",
        num_labels=len(label2id),
        id2label={str(k): v for k, v in id2label.items()},
        label2id=label2id
    )

    # 4. 训练参数
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=10,  # 训练轮数
        per_device_train_batch_size=16,  # 批次大小
        per_device_eval_batch_size=16,
        learning_rate=2e-5,  # 学习率
        weight_decay=0.01,  # 权重衰减
        warmup_ratio=0.1,  # 预热比例
        logging_dir='./logs',
        logging_steps=50,
        eval_strategy="epoch",  # 每个epoch评估一次
        save_strategy="epoch",  # 每个epoch保存一次
        load_best_model_at_end=True,  # 结束时加载最佳模型
        metric_for_best_model="f1",
        greater_is_better=True,
        save_total_limit=3,  # 只保留最后3个模型
        fp16=False,  # 没有GPU设为False
        report_to="none",  # 不报告到wandb等
    )

    # 5. 数据整理器
    data_collator = DataCollatorForTokenClassification(tokenizer)

    # 6. 创建训练器
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics
    )

    # 7. 开始训练
    print("\n开始训练...")
    trainer.train()

    # 8. 保存模型
    print(f"\n保存模型到: {output_dir}")
    trainer.save_model()
    tokenizer.save_pretrained(output_dir)

    # 9. 保存标签映射
    with open(os.path.join(output_dir, 'label2id.json'), 'w', encoding='utf-8') as f:
        json.dump(label2id, f, indent=2, ensure_ascii=False)

    with open(os.path.join(output_dir, 'id2label.json'), 'w', encoding='utf-8') as f:
        json.dump(id2label, f, indent=2, ensure_ascii=False)

    print("训练完成!")
    return trainer


if __name__ == "__main__":
    # 加载标签映射
    with open('data/out/labels_info.json', 'r', encoding='utf-8') as f:
        labels_info = json.load(f)
    label2id = labels_info['label2id']
    id2label = {int(k): v for k, v in labels_info['id2label'].items()}

    # 训练模型
    train_model(
        train_file='data/out/bert_training_data_ecg_augmented.json',  # 增强后的数据
        valid_file='data/out/bert_training_data.json',  # 原始数据作为验证集
        label2id=label2id,
        output_dir='./ecg_ner_model'
    )