🚀 实体BERT模型

boltuix/EntityBERT 模型是一个轻量级、经过微调的Transformer模型，专为英文文本的命名实体识别（Named Entity Recognition, NER）任务而设计。它基于 boltuix/bert-mini 基础模型构建，能够高效识别36种实体类型，如人物、组织、地点和日期等，适用于信息提取、聊天机器人和搜索增强等应用场景。

🚀 快速开始

🧪 推理代码

可以使用以下Python代码运行命名实体识别任务：

from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained("boltuix/EntityBERT")
model = AutoModelForTokenClassification.from_pretrained("boltuix/EntityBERT")

# 输入文本
text = "Elon Musk launched Tesla in California on March 2025."
inputs = tokenizer(text, return_tensors="pt")

# 运行推理
with torch.no_grad():
    outputs = model(**inputs)
predictions = outputs.logits.argmax(dim=-1)

# 将预测结果映射到标签
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
label_map = model.config.id2label
labels = [label_map[p.item()] for p in predictions[0]]

# 打印结果
for token, label in zip(tokens, labels):
    if token not in tokenizer.all_special_tokens:
        print(f"{token:15} → {label}")

✨ 示例输出

Elon            → B-PERSON
Musk            → I-PERSON
launched        → O
Tesla           → B-ORG
in              → O
California      → B-GPE
on              → O
March           → B-DATE
2025            → I-DATE
.               → O

🛠️ 环境要求

pip install transformers torch pandas pyarrow

• Python：3.8+
• 存储空间：模型权重约需15 MB
• 可选：使用 seqeval 进行评估，使用 cuda 进行GPU加速

✨ 主要特性

🌈 模型描述

boltuix/EntityBERT 模型是一个轻量级、经过微调的Transformer模型，用于命名实体识别（NER）。它基于 boltuix/bert-mini 基础模型构建，经过优化，能够高效识别英文文本中的36种实体类型（如人物、组织、地点、日期等），非常适合信息提取、聊天机器人和搜索增强等应用场景。

• 数据集：boltuix/conll2025-ner（143,709条记录，6.38 MB）
• 实体类型：36种NER标签（18种实体类别，带有B-/I-标签 + O）
• 训练示例：约115,812条 | 验证集：约15,680条 | 测试集：约12,217条
• 适用领域：新闻、用户生成内容、研究语料库
• 支持任务：句子级和文档级的NER任务
• 版本：v1.0

🔧 模型信息

• 开发者：Boltuix
• 许可证：Apache-2.0
• 支持语言：英文
• 模型类型：基于Transformer的分词分类模型
• 训练时间：2025年6月11日前
• 基础模型：boltuix/bert-mini
• 参数数量：约440万
• 模型大小：约15 MB

🔗 相关链接

• 模型仓库：boltuix/EntityBERT（占位符，需更新为正确的URL）
• 数据集：boltuix/conll2025-ner（占位符，需更新为正确的URL）
• Hugging Face文档：Transformers
• 演示：即将推出

🎯 命名实体识别的应用场景

🌟 直接应用

• 信息提取：从文章、博客或报告中识别姓名（👤 人物）、地点（🌍 地理政治实体）和日期（🗓️ 日期）。
• 聊天机器人和虚拟助手：通过识别实体来提高对用户查询的理解能力。
• 搜索增强：实现基于实体的语义搜索（例如，“2025年关于巴黎的新闻”）。
• 知识图谱：构建连接实体（如 🏢 组织和 👤 人物）的结构化图谱。

🌱 下游任务

• 领域适配：针对医学 🩺、法律 📜 或金融 💸 等专业领域进行微调。
• 多语言扩展：针对非英文语言进行重新训练。
• 自定义实体：适应特定领域（例如，产品ID、股票代码）。

❌ 局限性

• 仅支持英文：默认情况下仅适用于英文文本。
• 领域偏差：基于 boltuix/conll2025-ner 数据集进行训练，可能更适用于新闻和正式文本，在非正式或社交媒体内容上的表现可能较弱。
• 泛化能力：对于数据集中未包含的罕见或高度上下文相关的实体，可能难以准确识别。

📦 安装指南

pip install transformers torch pandas pyarrow seqeval

• Python：3.8+
• 存储空间：模型约需15 MB，数据集约需6.38 MB
• 可选：使用NVIDIA CUDA进行GPU加速

下载说明 📥

• 模型：boltuix/EntityBERT

🧠 实体标签

该模型支持 boltuix/conll2025-ner 数据集中的36种NER标签，采用BIO标签方案：

• B-：实体的开始
• I-：实体的内部
• O：非实体

示例： 文本："Tesla opened in Shanghai on April 2025" 标签：[B-ORG, O, O, B-GPE, O, B-DATE, I-DATE]

📈 性能表现

使用 seqeval 在 boltuix/conll2025-ner 测试集（约12,217个示例）上进行评估：

注意：在不同领域或文本类型上，性能可能会有所不同。

🔧 技术细节

⚙️ 训练设置

• 硬件：NVIDIA GPU
• 训练时间：约1.5小时
• 参数数量：约440万
• 优化器：AdamW
• 精度：FP32
• 批量大小：16
• 学习率：2e-5

🧠 模型训练

可以在 boltuix/conll2025-ner 数据集上对 boltuix/bert-mini 进行微调，以复现或扩展 EntityBERT 模型。以下是一个简化的训练脚本：

# 🛠️ 步骤1：安静地安装所需库
!pip install evaluate transformers datasets tokenizers seqeval pandas pyarrow -q

# 🚫 步骤2：禁用Weights & Biases (WandB)
import os
os.environ["WANDB_MODE"] = "disabled"

# 📚 步骤2：导入必要的库
import pandas as pd
import datasets
import numpy as np
from transformers import BertTokenizerFast
from transformers import DataCollatorForTokenClassification
from transformers import AutoModelForTokenClassification
from transformers import TrainingArguments, Trainer
import evaluate
from transformers import pipeline
from collections import defaultdict
import json

# 📥 步骤3：从Parquet文件加载CoNLL - 2025 NER数据集
# 下载地址：https://huggingface.co/datasets/boltuix/conll2025-ner/blob/main/conll2025_ner.parquet
parquet_file = "conll2025_ner.parquet"
df = pd.read_parquet(parquet_file)

# 🔍 步骤4：将pandas DataFrame转换为Hugging Face数据集
conll2025 = datasets.Dataset.from_pandas(df)

# 🔎 步骤5：检查数据集结构
print("数据集结构:", conll2025)
print("数据集特征:", conll2025.features)
print("第一个示例:", conll2025[0])

# 🏷️ 步骤6：提取唯一标签并创建映射
# 由于ner_tags是字符串，收集所有唯一标签
all_tags = set()
for example in conll2025:
    all_tags.update(example["ner_tags"])
unique_tags = sorted(list(all_tags))  # 排序以保持一致性
num_tags = len(unique_tags)
tag2id = {tag: i for i, tag in enumerate(unique_tags)}
id2tag = {i: tag for i, tag in enumerate(unique_tags)}
print("唯一标签的数量:", num_tags)
print("唯一标签:", unique_tags)

# 🔧 步骤7：将字符串ner_tags转换为索引
def convert_tags_to_ids(example):
    example["ner_tags"] = [tag2id[tag] for tag in example["ner_tags"]]
    return example

conll2025 = conll2025.map(convert_tags_to_ids)

# 📊 步骤8：根据'split'列拆分数据集
dataset_dict = {
    "train": conll2025.filter(lambda x: x["split"] == "train"),
    "validation": conll2025.filter(lambda x: x["split"] == "validation"),
    "test": conll2025.filter(lambda x: x["split"] == "test")
}
conll2025 = datasets.DatasetDict(dataset_dict)
print("拆分后的数据集结构:", conll2025)

# 🪙 步骤9：初始化分词器
tokenizer = BertTokenizerFast.from_pretrained("boltuix/bert-mini")

# 📝 步骤10：对示例文本进行分词并检查
example_text = conll2025["train"][0]
tokenized_input = tokenizer(example_text["tokens"], is_split_into_words=True)
tokens = tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])
word_ids = tokenized_input.word_ids()
print("单词ID:", word_ids)
print("分词后的输入:", tokenized_input)
print("ner_tags与输入ID的长度:", len(example_text["ner_tags"]), len(tokenized_input["input_ids"]))

# 🔄 步骤11：定义分词和对齐标签的函数
def tokenize_and_align_labels(examples, label_all_tokens=True):
    """
    对输入进行分词并对齐NER任务的标签。

    参数:
        examples (dict): 包含tokens和ner_tags的字典。
        label_all_tokens (bool): 是否对所有子词标记进行标签。

    返回:
        dict: 带有对齐标签的分词输入。
    """
    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)  # 特殊标记设为 -100
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])  # 单词的第一个标记获取标签
            else:
                label_ids.append(label[word_idx] if label_all_tokens else -100)  # 子词获取标签或 -100
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

# 🧪 步骤12：测试分词和标签对齐
q = tokenize_and_align_labels(conll2025["train"][0:1])
print("分词和对齐后的示例:", q)

# 📋 步骤13：打印标记及其对应的标签
for token, label in zip(tokenizer.convert_ids_to_tokens(q["input_ids"][0]), q["labels"][0]):
    print(f"{token:_<40} {label}")

# 🔧 步骤14：对整个数据集应用分词
tokenized_datasets = conll2025.map(tokenize_and_align_labels, batched=True)

# 🤖 步骤15：使用正确的标签数量初始化模型
model = AutoModelForTokenClassification.from_pretrained("boltuix/bert-mini", num_labels=num_tags)

# ⚙️ 步骤16：设置训练参数
args = TrainingArguments(
    "boltuix/bert-ner",
    eval_strategy="epoch", # 将evaluation_strategy改为eval_strategy
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=1,
    weight_decay=0.01,
    report_to="none"
)
# 📊 步骤17：初始化数据收集器以进行动态填充
data_collator = DataCollatorForTokenClassification(tokenizer)

# 📈 步骤18：加载评估指标
metric = evaluate.load("seqeval")

# 🏷️ 步骤19：设置标签列表并测试指标计算
label_list = unique_tags
print("标签列表:", label_list)
example = conll2025["train"][0]
labels = [label_list[i] for i in example["ner_tags"]]
print("指标测试:", metric.compute(predictions=[labels], references=[labels]))

# 📉 步骤20：定义计算评估指标的函数
def compute_metrics(eval_preds):
    """
    计算NER的精确率、召回率、F1和准确率。

    参数:
        eval_preds (tuple): 预测的logits和真实标签。

    返回:
        dict: 评估指标。
    """
    pred_logits, labels = eval_preds
    pred_logits = np.argmax(pred_logits, axis=2)
    predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(pred_logits, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(pred_logits, labels)
    ]
    results = metric.compute(predictions=predictions, references=true_labels)
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"],
    }

# 🚀 步骤21：初始化并训练训练器
trainer = Trainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)
trainer.train()

# 💾 步骤22：保存微调后的模型
model.save_pretrained("boltuix/bert-ner")
tokenizer.save_pretrained("tokenizer")

# 🔗 步骤23：使用标签映射更新模型配置
id2label = {str(i): label for i, label in enumerate(label_list)}
label2id = {label: str(i) for i, label in enumerate(label_list)}
config = json.load(open("boltuix/bert-ner/config.json"))
config["id2label"] = id2label
config["label2id"] = label2id
json.dump(config, open("boltuix/bert-ner/config.json", "w"))

# 🔄 步骤24：加载微调后的模型
model_fine_tuned = AutoModelForTokenClassification.from_pretrained("boltuix/bert-ner")

# 🛠️ 步骤25：创建用于NER推理的管道
nlp = pipeline("token-classification", model=model_fine_tuned, tokenizer=tokenizer)

# 📝 步骤26：对示例句子进行NER
example = "On July 4th, 2023, President Joe Biden visited the United Nations headquarters in New York to deliver a speech about international law and donated $5 million to relief efforts."
ner_results = nlp(example)
print("第一个示例的NER结果:", ner_results)

# 📍 步骤27：对房产地址进行NER并格式化输出
example = "This page contains information about the property located at 1275 Kinnear Rd, Columbus, OH, 43212."
ner_results = nlp(example)

# 🧹 步骤28：将NER结果处理为结构化实体
entities = defaultdict(list)
current_entity = ""
current_type = ""

for item in ner_results:
    entity = item["entity"]
    word = item["word"]
    if word.startswith("##"):
        current_entity += word[2:]  # 处理子词标记
    elif entity.startswith("B-"):
        if current_entity and current_type:
            entities[current_type].append(current_entity.strip())
        current_type = entity[2:].lower()
        current_entity = word
    elif entity.startswith("I-") and entity[2:].lower() == current_type:
        current_entity += " " + word  # 继续同一实体
    else:
        if current_entity and current_type:
            entities[current_type].append(current_entity.strip())
        current_entity = ""
        current_type = ""

# 如果存在最终实体，则添加
if current_entity and current_type:
    entities[current_type].append(current_entity.strip())

# 📤 步骤29：输出最终的JSON
final_json = dict(entities)
print("结构化的NER输出:")
print(json.dumps(final_json, indent=2))

🛠️ 提示

• 超参数：可以尝试调整 learning_rate（1e-5 到 5e-5）或 num_train_epochs（2 - 5）。
• GPU：使用 fp16=True 可以加快训练速度。
• 自定义数据：可以修改脚本以适应自定义的NER数据集。

⏱️ 预期训练时间

在NVIDIA GPU（例如，T4）上，对于约115,812个示例，进行3个epoch，批量大小为16，大约需要1.5小时。

🌍 碳影响

• 排放量：约40克CO₂eq（通过ML Impact工具估算，在GPU上训练1.5小时）

📄 许可证

本项目采用Apache-2.0许可证。