daft.functions.llm_generate#

def llm_generate(
    text: Expression,
    model: str = "facebook/opt-125m",
    provider: Literal["vllm", "openai"] = "vllm",
    concurrency: int = 1,
    batch_size: int | None = None,
    num_cpus: int | None = None,
    num_gpus: int | None = None,
    **generation_config: dict[str, Any],
) -> Expression:
    """A UDF for running LLM inference over an input column of strings.

    This UDF provides a flexible interface for text generation using various LLM providers.
    By default, it uses vLLM for efficient local inference.

    Args:
        text (String Expression):
            The input text column to generate from
        model (str, default="facebook/opt-125m"):
            The model identifier to use for generation
        provider (str, default="vllm"):
            The LLM provider to use for generation. Supported values: "vllm", "openai"
        concurrency (int, default=1):
            The number of concurrent instances of the model to run
        batch_size (int, default=None):
            The batch size for the UDF. If None, the batch size will be determined by defaults based on the provider.
        num_cpus (int, default=None):
            The number of CPUs to use for the UDF
        num_gpus (int, default=None):
            The number of GPUs to use for the UDF
        generation_config (dict, default={}):
            Configuration parameters for text generation (e.g., temperature, max_tokens)

    Returns:
        Expression (String Expression): The generated text column

    Examples:
        Use vLLM provider:
        >>> import daft
        >>> from daft import col
        >>> from daft.functions import llm_generate, format
        >>>
        >>> df = daft.from_pydict({"city": ["Paris", "Tokyo", "New York"]})
        >>> df = df.with_column(
        ...     "description",
        ...     llm_generate(
        ...         format(
        ...             "Describe the main attractions and unique features of this city: {}.",
        ...             col("city"),
        ...         ),
        ...         model="facebook/opt-125m",
        ...     ),
        ... )
        >>> df.collect()

        Use OpenAI provider:
        >>> df = daft.from_pydict({"city": ["Paris", "Tokyo", "New York"]})
        >>> df = df.with_column(
        ...     "description",
        ...     llm_generate(
        ...         format(
        ...             "Describe the main attractions and unique features of this city: {}.",
        ...             col("city"),
        ...         ),
        ...         model="gpt-4o",
        ...         api_key="xxx",
        ...         provider="openai",
        ...     ),
        ... )
        >>> df.collect()

    Note:
        Make sure the required provider packages are installed (e.g. vllm, transformers, openai).
    """
    warnings.warn(
        "This method is deprecated and will be removed in v0.8.0. Use daft.functions.prompt instead.",
        DeprecationWarning,
        stacklevel=2,
    )

    cls: Any = None
    if provider == "vllm":
        cls = _vLLMGenerator
        if batch_size is None:
            batch_size = 1024
    elif provider == "openai":
        cls = _OpenAIGenerator
        if batch_size is None:
            batch_size = 128
    else:
        raise ValueError(f"Unsupported provider: {provider}")

    llm_generator = udf(
        return_dtype=DataType.string(),
        batch_size=batch_size,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
        concurrency=concurrency,
    )(cls).with_init_args(
        model=model,
        generation_config=generation_config,
    )

    return llm_generator(text)