EmbeddingGeneration

使用 Embeddings 模型生成 embeddings。

EmbeddingGeneration 是一个 Step，它使用 Embeddings 模型为提供的输入文本生成句子 embeddings。

属性

• embeddings: 用于生成句子 embeddings 的 Embeddings 模型。

输入和输出列

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

    subgraph EmbeddingGeneration
        StepInput[Input Columns: text]
        StepOutput[Output Columns: embedding]
    end

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

输入

• text (str): 需要为其生成句子 embedding 的文本。

输出

• embedding (List[Union[float, int]]): 生成的句子 embedding。

示例

使用 Sentence Transformers 生成句子 embeddings

from distilabel.models import SentenceTransformerEmbeddings
from distilabel.steps import EmbeddingGeneration

embedding_generation = EmbeddingGeneration(
    embeddings=SentenceTransformerEmbeddings(
        model="mixedbread-ai/mxbai-embed-large-v1",
    )
)

embedding_generation.load()

result = next(embedding_generation.process([{"text": "Hello, how are you?"}]))
# [{'text': 'Hello, how are you?', 'embedding': [0.06209656596183777, -0.015797119587659836, ...]}]