pyflink.dataframe.ai.llm.LLMAccessor.ai_embed#

LLMAccessor.ai_embed(input_col: Union[str, Expression], dimension: int = 1024, *, provider: Optional[str] = None, model: Optional[str] = None, cache_table: Optional[str] = None, cache_key: Optional[Union[str, List[str]]] = None, config: Optional[Dict[str, str]] = None, **kwargs) → DataFrame[source]#

Generate embedding vectors for text.

Parameters

input_col – Column name (str) or Expression for the input text.
dimension – Dimension of the embedding vector (default 1024).
provider – Registered model provider lookup name. Uses default if not specified. If no provider is configured, model is treated as a catalog model name.
model – Model name.
cache_table – Pre-registered catalog table identifier for the embedding cache.
cache_key – Optional cache key column name or list of column names.
config – Optional dict of runtime prediction options.
**kwargs – Per-call overrides for the selected DataFrame model provider’s constructor parameters. Overrides are applied only when the call uses a configured provider.

Returns

A new DataFrame with a column appended –

embedding (ARRAY<FLOAT>): the embedding vector.

Examples

Simple invocation:

::

>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
...     task="embeddings"))
>>> df = pf.from_dict({
...     "tenant_id": ["tenant-a"],
...     "text": ["Hello world"],
... })
>>> df.llm.ai_embed("text", 512, model="text-embedding-v4")

With cache:

::

>>> df.llm.ai_embed(
...     "text", 512, model="text-embedding-v4",
...     cache_table="`fluss-catalog`.default_database.embed_cache",
...     cache_key="text")

Using runtime config supported by pyflink.dataframe.OpenAICompatProvider:

::

>>> df.llm.ai_embed(
...     "text", 512, model="text-embedding-v4",
...     config={"max-concurrent-operations": "20"})