EmbeddingTaskGenerator

使用 LLM 为 embedding 相关任务生成任务描述。

EmbeddingTaskGenerator 是一个 GeneratorTask，除了提供的属性外，不接收任何输入，它使用基于 category 属性的预定义提示，为 embedding 相关任务生成任务描述。 category 属性应该是以下之一

- `text-retrieval`: Generate task descriptions for text retrieval tasks.
- `text-matching-short`: Generate task descriptions for short text matching tasks.
- `text-matching-long`: Generate task descriptions for long text matching tasks.
- `text-classification`: Generate task descriptions for text classification tasks.

属性

• category：要生成的任务的类别，可以是 text-retrieval、text-matching-short、text-matching-long 或 text-classification。

• flatten_tasks：是否展平任务，即，由于 LLM 生成任务列表，此属性指示是否展平列表。默认为 False，这意味着使用 num_generations=1 运行此任务将返回一个 distilabel.Distiset，其中只有一行包含大约 20 个任务的列表；否则，如果设置为 True，它将返回一个 distilabel.Distiset，其中包含大约 20 行，每行包含一个任务。

输入 & 输出列

graph TD
    subgraph Dataset
        subgraph New columns
            OCOL0[tasks]
            OCOL1[task]
            OCOL2[model_name]
        end
    end

    subgraph EmbeddingTaskGenerator
        StepOutput[Output Columns: tasks, task, model_name]
    end

    StepOutput --> OCOL0
    StepOutput --> OCOL1
    StepOutput --> OCOL2

输出

• tasks (List[str])：LLM 生成的任务列表。

• task (str)：如果 flatten_tasks=True，则为 LLM 生成的任务。

• model_name (str)：用于生成任务的模型的名称。

示例

为文本检索生成 embedding 任务

from distilabel.pipeline import Pipeline
from distilabel.steps.tasks import EmbeddingTaskGenerator

with Pipeline("my-pipeline") as pipeline:
    task = EmbeddingTaskGenerator(
        category="text-retrieval",
        flatten_tasks=True,
        llm=...,  # LLM instance
    )

    ...

    task >> ...

参考文献

• 使用大型语言模型改进文本 Embedding