跳到内容

GenerateEmbeddings

使用 LLM 的最后一个隐藏状态生成 embeddings。

使用 LLM 的最后一个隐藏状态为文本输入生成 embeddings,如论文 'What Makes Good Data for Alignment? A Comprehensive Study of Automatic Data Selection in Instruction Tuning' 中所述。

属性

  • llm: 用于生成 embeddings 的 LLM

输入和输出列

graph TD
    subgraph Dataset
        subgraph Columns
            ICOL0[text]
        end
        subgraph New columns
            OCOL0[embedding]
            OCOL1[model_name]
        end
    end

    subgraph GenerateEmbeddings
        StepInput[Input Columns: text]
        StepOutput[Output Columns: embedding, model_name]
    end

    ICOL0 --> StepInput
    StepOutput --> OCOL0
    StepOutput --> OCOL1
    StepInput --> StepOutput

输入

  • text (str, List[Dict[str, str]]): 用于生成 embeddings 的输入文本或对话。

输出

  • embedding (List[float]): 输入文本或对话的 embedding。

  • model_name (str): 用于生成 embeddings 的模型名称。

示例

对 LLM 候选进行排序

from distilabel.steps.tasks import GenerateEmbeddings
from distilabel.models.llms.huggingface import TransformersLLM

# Consider this as a placeholder for your actual LLM.
embedder = GenerateEmbeddings(
    llm=TransformersLLM(
        model="TaylorAI/bge-micro-v2",
        model_kwargs={"is_decoder": True},
        cuda_devices=[],
    )
)
embedder.load()

result = next(
    embedder.process(
        [
            {"text": "Hello, how are you?"},
        ]
    )
)

参考