xpark.dataset.ImageCaption#

class xpark.dataset.ImageCaption(_local_model: str | None = 'Qwen/Qwen3.5-4B', /, *, prompt_builder: Callable[[PromptComponents], str] | None = None, max_words: int = 100, quality: Literal['fast', 'balanced', 'accurate'] | None = None, min_token_num: int | None = None, max_token_num: int | None = None, fallback_response: str | None = None, sampling_params: Mapping[str, Any] | None = None, chat_template_kwargs: Mapping[str, Any] | None = None, engine_kwargs: Mapping[str, Any] | None = None, base_url: str | None = None, model: str | None = None, api_key: str = 'NOT_SET', max_qps: int | None = None, max_concurrency: int | None = None, max_retries: int | None = None)#

Image caption processor for CPU, GPU and remote Http requests.

Generates a natural-language description (caption) for each input image using a Qwen3.5 vision-language model. The backend is chosen automatically per worker: GPU (in-process vLLM) when a CUDA device is assigned, remote OpenAI-compatible endpoint on IO workers, otherwise CPU (GGUF + mmproj via xllamacpp).

Each entry in Args carries the backends it applies to: [all], [CPU/GPU] (local engines) or [remote only] (IO workers).

Parameters:

  • _local_model – [CPU/GPU] The vision model name for CPU or GPU. Available models: [‘Qwen/Qwen3.5-0.8B’, ‘Qwen/Qwen3.5-4B’, ‘Qwen/Qwen3.5-9B’]. Defaults to Qwen/Qwen3.5-4B.

  • prompt_builder – [all] Optional (components) -> str hook that replaces the whole system prompt. components is a PromptComponents: the blocks xpark would render (roles_and_tasks, response_format, default_prompt), the row’s prepared image and the token_budget left for text — a hard ceiling: a longer prompt raises ValueError instead of being truncated. Runs once per row, on the worker that owns the batch, so keep it stateless.

  • max_words – [all] A non-negative best-effort target for the number of words in the caption. Defaults to 100. If set to 0, there is no word limit. A soft, prompt-level constraint rendered into a <response_format> block, orthogonal to the hard token cap sampling_params["max_tokens"].

  • quality –

    [all] Semantic preset for the per-image vision-token budget. Sets max_token_num under the hood:

    • "fast" — max_token_num=256: thumbnails, coarse captioning, low-latency / high-throughput.

    • "balanced" — max_token_num=1280: general understanding, the usual choice.

    • "accurate" — max_token_num=16384: keeps fine text / small objects intact (OCR-grade).

    The pixel budget is max_token_num * factor**2, where factor is inferred from the model by smart_resize_factor_for_model() (32 for Qwen3-VL / Qwen3.5, 28 for the Qwen2 family).

    When None (default) no preset is applied and max_token_num decides. When combined with an explicit max_token_num, the tighter (smaller) of the two wins.

  • min_token_num – [all] lower bound on the per-image vision-token budget. Defaults to 128.

  • max_token_num – [all] upper bound on the per-image vision-token budget. Defaults to the 16384 the checkpoint’s own image processor allows (the resolved budget caps it lower on CPU). Combined with quality via min(...).

  • fallback_response – [all] The response value to return when the backend produces no caption text (remote request failure after retries, or empty local generation). If set to None, an empty string is returned instead of a null.

  • engine_kwargs – [CPU/GPU] Extra arguments forwarded to the local inference engine (GPU: vllm.LLM constructor; CPU: llama.cpp CommonParams fields). Unknown keys are dropped with a warning on GPU and raise on CPU. Overriding fields the backend sizes itself (e.g. n_ctx) is allowed but gives up the capacity-sizing guarantees. Ignored by the remote backend.

  • base_url – [remote only] The base URL of the remote OpenAI-compatible endpoint.

  • model – [remote only] The request model name.

  • api_key – [remote only] The request API key.

  • max_qps – [remote only] The maximum query-per-second rate for remote requests.

  • max_concurrency – [remote only] The maximum number of in-flight remote requests allowed concurrently.

  • max_retries – [remote only] The maximum number of retries per request, with exponential backoff. Defaults to 2.

  • sampling_params – [all] Sampling parameters for the sampler: on GPU they become SamplingParams, on CPU and on remote HTTP they go into the chat-completions request. Layered field by field over the model’s presets, so {"temperature": 0.2} keeps every other preset value. max_tokens is the hard generation cap; it defaults to 1024 when omitted. Unknown fields are dropped with a warning on GPU. repetition_penalty only takes effect on the GPU backend: the CPU shim does not expose it, and on the remote endpoint it travels via extra_body, so whether it applies is up to the server.

  • chat_template_kwargs – [all] Options for the chat template (apply_chat_template on GPU, chat_template_kwargs in the request on CPU / remote HTTP). Layered over xpark’s {"enable_thinking": False}, so pass {"enable_thinking": True} to let the model reason (at the cost of the generation budget).

Raises:

ValueError – If an argument is out of range or contradicts another (max_words, min_token_num / max_token_num, quality), the model name is unsupported, or the target itself is missing (neither _local_model nor base_url; _local_model on a local worker; base_url / model on an IO worker). Also raised when a budget cannot be satisfied (the clip-level token budget exceeds the model context, or the worst legal image does not fit the CPU engine window), when engine_kwargs names a field the CPU backend does not have, and when prompt_builder is not callable, returns a blank or non-string prompt, or returns a prompt longer than the token_budget the worker’s engine has left for text.

Examples

from xpark.dataset.expressions import col
from xpark.dataset import ImageCaption, from_items
from xpark.dataset.utils import PromptComponents, estimate_prompt_tokens, format_prompt

ds = from_items(["cat.jpg", "dog.jpg"])

# CPU worker.
ds = ds.with_column(
    "caption_cpu",
    ImageCaption().options(num_workers={"CPU": 1}).with_column(col("item")),
)

# GPU worker.
ds = ds.with_column(
    "caption_gpu",
    ImageCaption().options(num_workers={"GPU": 1}).with_column(col("item")),
)

# Remote OpenAI-compatible endpoint (IO worker).
ds = ds.with_column(
    "caption_io",
    ImageCaption(
        base_url="https://your-endpoint/v1",
        model="youtu-vita",
    )
    .options(num_workers={"IO": 1})
    .with_column(col("item")),
)

# --- Custom prompts ---------------------------------------------------------

# 1. Keep xpark's prompt and append one constraint.
def ocr_prompt(components: PromptComponents) -> str:
    return components.default_prompt + "\nTranscribe every visible string, in reading order."

# 2. Swap the task wording only: xpark's ``<response_format>`` block survives.
def curator_prompt(components: PromptComponents) -> str:
    return format_prompt(
        roles_and_tasks=(
            "You are a museum curator. Describe the artwork's subject, medium and period."
        ),
        response_format=components.response_format,
    )

# 3. Build the system turn from scratch: the blocks, the prepared image and the
#    budget are still yours to read. The user turn (the image item) stays xpark's.
def grounded_ocr_prompt(components: PromptComponents) -> str:
    height, width = components.image.shape[:2]
    lines = [
        "You are an OCR engine. Transcribe every string in the image, in reading order.",
        f"The image is {width}x{height} px.",
    ]
    prompt = "\n".join(lines)
    # ``token_budget`` is a hard ceiling: going over raises, so trim when close.
    if components.token_budget is not None:
        while estimate_prompt_tokens(prompt) > components.token_budget and len(lines) > 1:
            lines.pop()
            prompt = "\n".join(lines)
    return prompt

for name, builder in [
    ("caption_ocr", ocr_prompt),
    ("caption_curator", curator_prompt),
    ("caption_grounded", grounded_ocr_prompt),
]:
    ds = ds.with_column(
        name,
        ImageCaption(prompt_builder=builder)
        .options(num_workers={"CPU": 1})
        .with_column(col("item")),
    )

# Sampling and chat-template knobs: one destination each.
ds = ds.with_column(
    "caption_thinking",
    ImageCaption(
        sampling_params={"temperature": 0.2, "max_tokens": 512},
        chat_template_kwargs={"enable_thinking": True},
    )
    .options(num_workers={"CPU": 1})
    .with_column(col("item")),
)

print(ds.take(2))

Methods

__call__(images)

Call self as a function.

options(**kwargs)

with_column(images)

__call__(images: pa.ChunkedArray) → pa.Array#

Call self as a function.

options(**kwargs: Unpack[ExprUDFOptions]) → Self#
with_column(images: pa.ChunkedArray) → pa.Array#