SpaCy v2.x NER训练代码迁移至v3.x适配en_core_web_trf方法

以下是将spaCy v2.x版本的train_new_entity_type.py迁移到v3.x，且适配Transformer预训练模型en_core_web_trf的完整代码：

import spacy
from spacy.tokens import DocBin
from spacy.training import Example
import random
from pathlib import Path

LABEL = "ANIMAL"

# 训练数据（与v2.x示例一致）
TRAIN_DATA = [
    (
        "Horses are too tall and they pretend to care about your feelings",
        {"entities": [(0, 6, LABEL)]},
    ),
    ("Do they bite?", {"entities": []}),
    (
        "horses are too tall and they pretend to care about your feelings",
        {"entities": [(0, 6, LABEL)]},
    ),
    ("horses pretend to care about your feelings", {"entities": [(0, 6, LABEL)]}),
    (
        "they pretend to care about your feelings, those horses",
        {"entities": [(48, 54, LABEL)]},
    ),
    ("horses?", {"entities": [(0, 6, LABEL)]}),
]

def main(model="en_core_web_trf", new_model_name="animal_ner_trf", output_dir=None, n_iter=10):
    random.seed(0)
    # 加载预训练Transformer模型
    nlp = spacy.load(model)
    print(f"Loaded model '{model}'")

    # 获取NER组件并添加新标签
    ner = nlp.get_pipe("ner")
    ner.add_label(LABEL)
    # 可添加额外测试标签，不影响训练
    ner.add_label("VEGETABLE")

    # 将训练数据转换为spaCy v3.x要求的Example格式
    doc_bin = DocBin()
    for text, annotations in TRAIN_DATA:
        # 创建基础Doc对象
        doc = nlp.make_doc(text)
        # 构建训练用Example
        example = Example.from_dict(doc, annotations)
        doc_bin.add(example.reference)
    # 临时保存训练数据（也可直接在内存中处理）
    doc_bin.to_disk("./train_temp.spacy")

    # 配置训练参数：冻结Transformer相关组件，仅训练NER
    config = nlp.config.copy()
    # 设置针对NER微调的优化器（Transformer模型建议使用较小学习率）
    config["training"]["optimizer"] = {
        "@optimizers": "spacy.Adam.v1",
        "learn_rate": 1e-5,
        "beta1": 0.9,
        "beta2": 0.999
    }
    config["training"]["max_epochs"] = n_iter
    # 指定训练语料路径
    config["training"]["train_corpus"] = {
        "@corpora": "spacy.Corpus.v1",
        "path": "./train_temp.spacy",
        "max_length": 0
    }
    # 冻结除NER外的所有组件（避免破坏预训练Transformer特征）
    for pipe_name in nlp.pipe_names:
        if pipe_name != "ner":
            config["components"][pipe_name]["source"] = model
            config["components"][pipe_name]["frozen"] = True

    # 初始化训练器并开始训练
    trainer = spacy.TrainingLoop(config, nlp)
    print("Starting training...")
    final_losses = trainer.run()
    print(f"Training completed. Final losses: {final_losses}")

    # 测试训练后的模型
    test_text = "Do you like horses?"
    doc = nlp(test_text)
    print(f"\nEntities in '{test_text}':")
    for ent in doc.ents:
        print(f"- {ent.label_}: {ent.text}")

    # 保存模型到指定目录
    if output_dir is not None:
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True)
        nlp.meta["name"] = new_model_name
        nlp.to_disk(output_dir)
        print(f"\nSaved model to {output_dir}")

        # 验证保存的模型
        print(f"\nLoading saved model from {output_dir}...")
        nlp_saved = spacy.load(output_dir)
        doc_saved = nlp_saved(test_text)
        print(f"Entities in saved model for '{test_text}':")
        for ent in doc_saved.ents:
            print(f"- {ent.label_}: {ent.text}")

if __name__ == "__main__":
    # 可修改参数，比如指定输出目录：output_dir="./animal_ner_model"
    main(output_dir="./animal_ner_trf_model")

关键迁移说明 #

• 数据格式：v3.x使用Example对象和DocBin管理训练数据，替代原有的纯文本+字典格式
• 训练流程：使用TrainingLoop替代手动的nlp.update循环，配置更规范
• Transformer适配：冻结trf_tok2vec、trf_wordpiecer等Transformer相关组件，仅训练NER层，使用较小学习率避免破坏预训练特征
• 配置管理：基于原模型的config修改训练参数，更符合v3.x的设计规范

内容的提问来源于stack exchange，提问作者mCs