Token 分类¶

目标：展示 token 分类任务的标准工作流程，包括零样本建议和模型微调。
数据集：ontonotes5，一个大型语料库，包含需要为命名实体识别分类的各种文本类型。
库：datasets, gliner, transformers, spanmarker
组件：TextField, SpanQuestion, Suggestion, Query, Filter

开始使用¶

部署 Argilla 服务器¶

如果您已经部署了 Argilla，则可以跳过此步骤。否则，您可以按照本指南快速部署 Argilla。

设置环境¶

要完成本教程，您需要通过 pip 安装 Argilla SDK 和一些第三方库。

!pip install argilla

!pip install gliner==0.2.6 transformers==4.40.2 span_marker==1.5.0

让我们进行所需的导入

import re

import argilla as rg

import torch
from datasets import load_dataset, Dataset, DatasetDict
from gliner import GLiNER
from span_marker import SpanMarkerModel, Trainer
from transformers import TrainingArguments

您还需要使用 api_url 和 api_key 连接到 Argilla 服务器。

# Replace api_url with your url if using Docker
# Replace api_key with your API key under "My Settings" in the UI
# Uncomment the last line and set your HF_TOKEN if your space is private
client = rg.Argilla(
    api_url="https://[your-owner-name]-[your_space_name].hf.space",
    api_key="[your-api-key]",
    # headers={"Authorization": f"Bearer {HF_TOKEN}"}
)

数据集的氛围检查¶

我们将查看数据集，以了解其结构和包含的数据类型。我们通过使用嵌入式 Hugging Face Dataset Viewer 来做到这一点。

配置和创建 Argilla 数据集¶

现在，我们需要配置数据集。在设置中，我们可以指定指南、字段和问题。如果需要，您还可以添加元数据和向量。但是，对于我们的用例，我们只需要一个文本字段和一个跨度问题，分别对应于 token 和 tags 列。我们将专注于命名实体识别，但此工作流程也适用于跨度分类，后者的不同之处在于跨度不太明确定义且经常重叠。

labels = [
    "CARDINAL",
    "DATE",
    "PERSON",
    "NORP",
    "GPE",
    "LAW",
    "PERCENT",
    "ORDINAL",
    "MONEY",
    "WORK_OF_ART",
    "FAC",
    "TIME",
    "QUANTITY",
    "PRODUCT",
    "LANGUAGE",
    "ORG",
    "LOC",
    "EVENT",
]

settings = rg.Settings(
    guidelines="Classify individual tokens according to the specified categories, ensuring that any overlapping or nested entities are accurately captured.",
    fields=[
        rg.TextField(
            name="text",
            title="Text",
            use_markdown=False,
        ),
    ],
    questions=[
        rg.SpanQuestion(
            name="span_label",
            field="text",
            labels=labels,
            title="Classify the tokens according to the specified categories.",
            allow_overlapping=False,
        )
    ],
)

让我们使用名称和定义的设置创建数据集

dataset = rg.Dataset(
    name="token_classification_dataset",
    settings=settings,
)
dataset.create()

添加记录¶

我们已经创建了数据集（您可以在 UI 中查看它），但我们仍然需要添加用于标注的数据。在本例中，我们将使用来自 Hugging Face Hub 的 ontonote5 数据集。具体来说，我们将使用 test 拆分的 2100 个样本。

hf_dataset = load_dataset("tner/ontonotes5", split="test[:2100]")

我们将迭代 Hugging Face 数据集，将数据添加到 Argilla 数据集的 Record 对象中的相应字段。然后，我们将使用 log 轻松地将它们添加到数据集中。

records = [rg.Record(fields={"text": " ".join(row["tokens"])}) for row in hf_dataset]

dataset.records.log(records)

添加初始模型建议¶

下一步是向数据集添加建议。这将使标注团队的工作更轻松、更快速。建议将显示为预选选项，因此标注员只需更正它们即可。在我们的例子中，我们将使用 GLiNER 模型生成它们。但是，您可以使用您选择的框架或技术。

注意

有关更多信息，您可以查看 GLiNER 仓库和原始论文。

我们将首先加载预训练的 GLiNER 模型。具体来说，我们将使用 Hugging Face Hub 中提供的 gliner_mediumv2。

gliner_model = GLiNER.from_pretrained("urchade/gliner_mediumv2.1")

接下来，我们将创建一个函数来使用此通用模型生成预测，该模型可以识别指定的标签，而无需预先在这些标签上进行训练。该函数将返回一个字典，该字典使用必要的架构格式化，以将实体添加到我们的 Argilla 数据集中。此架构包括键“start”和“end”，以指示跨度开始和结束的索引，以及实体标签的“label”。

def predict_gliner(model, text, labels, threshold):
    entities = model.predict_entities(text, labels, threshold)
    return [
        {k: v for k, v in ent.items() if k not in {"score", "text"}} for ent in entities
    ]

要更新记录，我们需要从服务器检索它们并使用新建议更新它们。始终需要提供 id，因为它是记录的标识符，用于更新记录并避免创建新记录。

data = dataset.records.to_list(flatten=True)
updated_data = [
    {
        "span_label": predict_gliner(
            model=gliner_model, text=sample["text"], labels=labels, threshold=0.70
        ),
        "id": sample["id"],
    }
    for sample in data
]
dataset.records.log(records=updated_data)

瞧！我们已将建议添加到数据集中，它们将以 ✨ 标记在 UI 中显示。

使用 Argilla 评估¶

现在，我们可以开始标注过程。只需在 Argilla UI 中打开数据集并开始标注记录。如果建议正确，您只需单击“提交”。否则，您可以选择正确的标签。

注意

查看此操作指南，了解有关在 UI 中标注的更多信息。

训练您的模型¶

标注完成后，我们将拥有一个强大的数据集来训练我们的实体识别模型。对于我们的案例，我们将训练 SpanMarker 模型，但您可以选择任何您选择的模型。因此，让我们首先检索标注的记录。

注意

查看此操作指南，了解有关在 Argilla 中筛选和查询的更多信息。此外，您可以查看 Hugging Face 文档中关于微调 token 分类模型的信息。

dataset = client.datasets("token_classification_dataset")

在我们的案例中，我们使用批量视图提交了 2000 个标注。

status_filter = rg.Query(filter=rg.Filter(("response.status", "==", "submitted")))

submitted = dataset.records(status_filter).to_list(flatten=True)

SpanMarker 接受任何数据集，只要它具有 tokens 和 ner_tags 列。ner_tags 可以使用 IOB、IOB2、BIOES 或 BILOU 标注方案以及常规的非方案标签进行标注。在我们的案例中，我们选择使用 IOB 格式。因此，我们将定义一个函数来根据此方案提取标注的 NER 标签。

注意

有关更多信息，您可以查看 SpanMarker 文档。

def get_iob_tag_for_token(token_start, token_end, ner_spans):
    for span in ner_spans:
        if token_start &gt;= span["start"] and token_end &lt;= span["end"]:
            if token_start == span["start"]:
                return f"B-{span['label']}"
            else:
                return f"I-{span['label']}"
    return "O"


def extract_ner_tags(text, responses):
    tokens = re.split(r"(\s+)", text)
    ner_tags = []

    current_position = 0
    for token in tokens:
        if token.strip():
            token_start = current_position
            token_end = current_position + len(token)
            tag = get_iob_tag_for_token(token_start, token_end, responses)
            ner_tags.append(tag)
        current_position += len(token)

    return ner_tags

现在让我们提取它们并保存两个列表，其中包含 token 和 NER 标签，这将帮助我们构建数据集以训练 SpanMarker 模型。

tokens = []
ner_tags = []
for r in submitted:
    tags = extract_ner_tags(r["text"], r["span_label.responses"][0])
    tks = r["text"].split()
    tokens.append(tks)
    ner_tags.append(tags)

此外，我们将必须指示标签，并且它们应格式化为整数。因此，我们将检索它们并进行映射。

labels = list(set([item for sublist in ner_tags for item in sublist]))

id2label = {i: label for i, label in enumerate(labels)}
label2id = {label: id_ for id_, label in id2label.items()}

mapped_ner_tags = [[label2id[label] for label in ner_tag] for ner_tag in ner_tags]

最后，我们将创建一个包含训练集和验证集的数据集。

records = [
    {
        "tokens": token,
        "ner_tags": ner_tag,
    }
    for token, ner_tag in zip(tokens, mapped_ner_tags)
]
span_dataset = DatasetDict(
    {
        "train": Dataset.from_list(records[:1500]),
        "validation": Dataset.from_list(records[1501:2000]),
    }
)

现在，让我们准备训练我们的模型。为此，建议使用 GPU。您可以检查它是否可用，如下所示。

if torch.cuda.is_available():
    device = torch.device("cuda")
    print(f"Using {torch.cuda.get_device_name(0)}")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
    print("Using MPS device")
else:
    device = torch.device("cpu")
    print("No GPU available, using CPU instead.")

我们将定义我们的模型和参数。在本例中，我们将使用 Hugging Face Hub 中提供的 bert-base-cased，但也可以应用其他模型。

注意

训练参数继承自 Transformers 库。您可以在此处查看更多信息。

encoder_id = "bert-base-cased"
model = SpanMarkerModel.from_pretrained(
    encoder_id,
    labels=labels,
    model_max_length=256,
    entity_max_length=8,
)

args = TrainingArguments(
    output_dir="models/span-marker",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=1,
    weight_decay=0.01,
    warmup_ratio=0.1,
    fp16=False,  # Set to True if available
    logging_first_step=True,
    logging_steps=50,
    evaluation_strategy="steps",
    save_strategy="steps",
    eval_steps=500,
    save_total_limit=2,
    dataloader_num_workers=2,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=span_dataset["train"],
    eval_dataset=span_dataset["validation"],
)

让我们训练它！这一次，我们使用高质量的人工标注训练集，因此预期结果会有所提高。

trainer.train()

trainer.evaluate()

您可以将其本地保存或推送到 Hub。然后从那里加载它。

# Save and load locally
# model.save_pretrained("token_classification_model")
# model = SpanMarkerModel.from_pretrained("token_classification_model")

# Push and load in HF
# model.push_to_hub("[username]/token_classification_model")
# model = SpanMarkerModel.from_pretrained("[username]/token_classification_model")

是时候进行预测了！我们将设置一个函数，该函数使用 predict 方法来获取建议的标签。该模型将根据文本推断标签。该函数将以 Argilla 数据集的相应结构返回跨度。

def predict_spanmarker(model, text):
    entities = model.predict(text)
    return [
        {
            "start": ent["char_start_index"],
            "end": ent["char_end_index"],
            "label": ent["label"],
        }
        for ent in entities
    ]

由于训练数据质量更高，我们可以期望模型更好。因此，我们可以使用新模型的建议更新剩余的未标注记录。

data = dataset.records.to_list(flatten=True)
updated_data = [
    {
        "span_label": predict_spanmarker(model=model, text=sample["text"]),
        "id": sample["id"],
    }
    for sample in data
]
dataset.records.log(records=updated_data)

结论¶

在本教程中，我们展示了 token 分类任务的端到端示例。这可以作为基础，但它可以迭代执行并无缝集成到您的工作流程中，以确保高质量的数据管理和改进的结果。

我们首先配置数据集，添加记录，并根据 GLiNer 预测添加建议。在标注过程之后，我们使用标注的数据训练了 SpanMarker 模型，并使用新建议更新了剩余的记录。