Skip to content

EmbedJob

Adrian edited this page Oct 16, 2026 · 1 revision

EmbedJob

Text embeddings: one list[float] per row (array<float> on Spark).

dib.EmbedJob(
    model,                # spec or URI
    text,                 # "{column}" template, e.g. "{body}" or "{title}\n{body}"
    dimensions=None,      # request shorter vectors (engines that support it)
    normalize=False,      # L2-normalize vectors (any engine)
    engine="auto",
    engine_options={},
)

Input text

text uses the same template syntax as GenerateJob. A single column is "{col}". Instruction prefixes that some models expect go into the template, e.g. "query: {question}" (E5) or "Represent this sentence for searching relevant passages: {q}" (BGE queries).

Rows whose rendered text is empty or whitespace get InvalidInputError: text to embed is empty and are not sent to the model.

Options

  • dimensions: shorter embeddings for models trained for it (OpenAI text-embedding-3-*, Matryoshka models in vLLM). Supported by the openai, vllm and fake engines. With engine="transformers", or "auto" when it could pick transformers, the job is rejected at creation: set engine= explicitly.
  • normalize: L2-normalize after the engine. Many sentence-transformers models do it themselves (a Normalize module in the export), and then this is a no-op.

Every vector is checked: non-empty, finite numbers, and the requested dimension count. Violations are row errors (EngineError).

Engines

  • openai: one HTTP request per embed_max_inputs texts (default 256). If a request fails with a client error (one bad input fails the whole request), its texts are retried one by one so only the bad row fails. Token usage is per request, so prompt_tokens is filled only for single-text requests.
  • vllm: pooling-mode LLM, with dimensions via PoolingParams.
  • transformers: AutoModel + pooling. sentence-transformers exports are read without needing that package: the pooling mode (cls/mean/max/lasttoken), Normalize and max_seq_length. Exports with other modules (e.g. Dense) are rejected. Override with engine_options={"pooling": "cls"}.

Examples

# Databricks Foundation Model API
emb = dib.EmbedJob(model=dib.models.databricks("databricks-gte-large-en"), text="{chunk}")

# Self-hosted from a Volume, copied to local NVMe on each node
emb = dib.EmbedJob(
    model=dib.models.path("/Volumes/main/ml/weights/bge-m3", copy_to_local=True),
    text="{chunk}",
    normalize=True,
)
vectors = emb.run(chunks_df, backend="spark", output_col="embedding")
vectors.write.saveAsTable("main.rag.chunk_embeddings")

The array<float> column can feed Databricks Vector Search (Delta Sync with self-managed embeddings) directly.

Clone this wiki locally