文本分类¶

目标: 展示文本分类任务的标准工作流程，包括零样本建议和模型微调。
数据集: IMDB，一个电影评论数据集，需要分类为正面或负面。
库: datasets, transformers, setfit
组件: TextField, LabelQuestion, Suggestion, Query, Filter

开始入门¶

部署 Argilla 服务器¶

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

设置环境¶

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

!pip install argilla

!pip install setfit==1.0.3 transformers==4.40.2

让我们进行所需的导入

import argilla as rg

from datasets import load_dataset, Dataset
from setfit import SetFitModel, Trainer, get_templated_dataset, sample_dataset

您还需要使用 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 数据集¶

现在，我们将需要配置数据集。在设置中，我们可以指定指南、字段和问题。如果需要，您还可以添加元数据和向量。但是，对于我们的用例，我们只需要一个文本字段和一个标签问题，分别对应于 text 和 label 列。

注意

查看此操作指南以了解有关配置和创建数据集的更多信息。

labels = ["positive", "negative"]

settings = rg.Settings(
    guidelines="Classify the reviews as positive or negative.",
    fields=[
        rg.TextField(
            name="review",
            title="Text from the review",
            use_markdown=False,
        ),
    ],
    questions=[
        rg.LabelQuestion(
            name="sentiment_label",
            title="In which category does this article fit?",
            labels=labels,
        )
    ],
)

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

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

添加记录¶

即使我们已经创建了数据集，它仍然缺少要标注的信息（您可以在 UI 中查看）。我们将使用来自 Hugging Face Hub 的 imdb 数据集。具体来说，我们将使用 train 拆分的 100 个样本。

hf_dataset = load_dataset("imdb", split="train[:100]")

我们将使用 log 和映射轻松地将它们添加到数据集中，在映射中，我们指示列 text 是应添加到字段 review 的数据。

dataset.records.log(records=hf_dataset, mapping={"text": "review"})

添加初始模型建议¶

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

我们将首先定义一个示例训练集，其中包含所需的标签：positive 和 negative。使用 get_templated_dataset 将从默认模板创建句子：“This sentence is {label}.”

zero_ds = get_templated_dataset(
    candidate_labels=labels,
    sample_size=8,
)

现在，我们将准备一个函数来训练 SetFit 模型。

注意

有关进一步的自定义，您可以查看 SetFit 文档。

def train_model(model_name, dataset):
    model = SetFitModel.from_pretrained(model_name)

    trainer = Trainer(
        model=model,
        train_dataset=dataset,
    )

    trainer.train()

    return model

让我们训练模型。我们将使用 TaylorAI/bge-micro-v2，它在 Hugging Face Hub 中可用。

model = train_model(model_name="TaylorAI/bge-micro-v2", dataset=zero_ds)

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

# Save and load locally
# model.save_pretrained("text_classification_model")
# model = SetFitModel.from_pretrained("text_classification_model")

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

是时候进行预测了！我们将设置一个函数，该函数使用 predict 方法来获取建议的标签。模型将根据文本推断标签。

def predict(model, input, labels):
    model.labels = labels

    prediction = model.predict([input])

    return prediction[0]

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

data = dataset.records.to_list(flatten=True)
updated_data = [
    {
        "sentiment_label": predict(model, sample["review"], labels),
        "id": sample["id"],
    }
    for sample in data
]
dataset.records.log(records=updated_data)

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

使用 Argilla 评估¶

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

注意

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

训练您的模型¶

标注完成后，我们将拥有一个强大的数据集来训练主模型。在我们的例子中，我们将使用 SetFit 进行微调。但是，您可以选择最适合您要求的模型。因此，让我们首先检索已标注的记录。

注意

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

dataset = client.datasets("text_classification_dataset")

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

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

由于每个记录只有一个响应，我们可以直接检索选定的标签，并创建每个标签包含 8 个样本的训练集。我们为每个标签选择了 8 个样本，以获得用于少样本学习的平衡数据集。

train_records = [
    {
        "text": r["review"],
        "label": r["sentiment_label.responses"][0],
    }
    for r in submitted
]
train_dataset = Dataset.from_list(train_records)
train_dataset = sample_dataset(train_dataset, label_column="label", num_samples=8)

我们可以使用我们之前的函数来训练模型，但这次使用高质量的人工标注训练集。

model = train_model(model_name="TaylorAI/bge-micro-v2", dataset=train_dataset)

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

data = dataset.records.to_list(flatten=True)
updated_data = [
    {
        "sentiment_label": predict(model, sample["review"], labels),
        "id": sample["id"],
    }
    for sample in data
]
dataset.records.log(records=updated_data)

结论¶

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

我们首先配置数据集，添加记录，并训练零样本 SetFit 模型（作为一个例子）以添加建议。在标注过程之后，我们使用标注数据训练了一个新模型，并使用新建议更新了剩余的记录。