xpark.dataset.VideoCaption#

class xpark.dataset.VideoCaption(_local_model: str | None = 'Qwen/Qwen3.5-4B', /, *, prompt_builder: Callable[[PromptComponents], str] | None = None, max_words: int = 100, num_frames: int | None = None, quality: Literal['fast', 'balanced', 'accurate'] | None = None, keyframes_only: bool = False, video_mode: Literal['auto', 'video', 'frames'] = 'auto', 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)#

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

Generates a natural-language description (caption) for each input video. 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). The input column holds video paths or URLs.

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

Every sampled frame is passed through smart_resize before it reaches a backend: the frames of one clip share a total vision-token budget, so their resolution shrinks as num_frames grows. The alignment factor and the clip-level bound come from the model’s registry entry (VisionModelSpec); in video_mode="video" this client-side resize is what enforces the clip-level budget as a hard constraint (the model’s own video processor resizes again on top of it).

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 token_budget left for text — a hard ceiling: a longer prompt raises ValueError instead of being truncated — and the row’s prepared media — frames as [(timestamp, frame), ...] in frames mode, each frame the clip’s own (H, W, 3) uint8 array; video (that same (T, H, W, 3) clip array) in video mode. Building on components.default_prompt keeps the built-in “one clip, chronological order” wording; assembling from scratch means restating it. Runs once per row, on the worker that owns the batch, so keep it stateless — see Examples:.

  • 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"].

  • num_frames – [all] Number of frames to uniformly sample per video. Defaults to 16 (or to the quality preset’s frame count when one is selected). An explicit value always wins over the preset.

  • quality – [all] Semantic preset bundling num_frames and max_token_num: "fast" (4 frames, 256 tokens/frame), "balanced" (8 frames, 768), "accurate" (16 frames, 768, the raw defaults). An explicit num_frames wins outright; an explicit max_token_num can only tighten the preset’s cap.

  • keyframes_only – [all] Only decode keyframes (I-frames) when sampling. Defaults to False.

  • video_mode –

    [all] How the sampled clip is handed to the model. Defaults to "auto", which picks "frames" on every backend.

    • "frames" – pass the frames as independent images with the timestamp markers rendered into the prompt text. Works on every backend and model, but forgoes the temporal merge (full num_frames vision-token cost).

    • "video" – pass the clip as one video item. The model’s video processor merges consecutive frames 2:1, halving vision-token cost, and dates each merged pair by the average of frames_indices / fps. Requires both the local vLLM backend and a model whose registry entry declares "video"; every model listed above does.

    A requested mode is always checked against the model’s declared modes, so "video" where the model or the backend does not support it raises ValueError rather than silently downgrading.

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

  • max_token_num – [all] upper bound on the per-frame vision-token budget, never above 768 (VLM_HARD_PER_FRAME_TOKEN_CAP). The effective cap is the clip’s vision budget (the model’s own clip bound, or a share of its context) divided across the clip’s merged frames; an explicit value can only tighten it further.

  • 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 (required on IO workers).

  • 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 256 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, num_frames / min_token_num / max_token_num, quality), if video_mode is not one of "auto" / "video" / "frames" or the model does not declare the requested mode, if video_mode="video" is requested on a backend other than the local vLLM one, if the model name is unsupported, or if 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 clip 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 VideoCaption, from_items
from xpark.dataset.utils import PromptComponents, format_prompt

ds = from_items(["https://example.com/clip.mp4"])

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

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

# Remote OpenAI-compatible endpoint (IO worker).
ds = ds.with_column(
    "caption_io",
    VideoCaption(
        None,
        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 scene_prompt(components: PromptComponents) -> str:
    return components.default_prompt + "\nList each scene change with its timestamp."

# 2. Build the system turn from scratch, in ``frames`` mode: xpark hands you the
#    sampled timeline -- ``[(timestamp_seconds, (H, W, 3) uint8 frame), ...]`` -- and
#    keeps the user turn (the frames, with their timestamp markers).
def timeline_prompt(components: PromptComponents) -> str:
    frames = components.frames
    stamps = ", ".join(f"{seconds:.1f}s" for seconds, _ in frames)
    # Your wording becomes the task block; xpark still renders ``response_format``
    # around it, so the tags stay the same as the built-in prompt.
    return format_prompt(
        roles_and_tasks=(
            f"You are shown {len(frames)} frames sampled in chronological order "
            f"from ONE video, at {stamps}. "
            f"Describe what happens between consecutive frames."
        ),
        response_format=components.response_format,
    )

# 3. ``video`` mode (vLLM only): the clip is one prepared ``(T, H, W, 3)`` tensor.
def whole_clip_prompt(components: PromptComponents) -> str:
    num_frames = components.video.shape[0]
    return (
        f"You are shown one {num_frames}-frame video clip. "
        f"Summarize it in a single sentence. {components.response_format}"
    )

for name, builder in [("caption_scenes", scene_prompt), ("caption_timeline", timeline_prompt)]:
    ds = ds.with_column(
        name,
        VideoCaption(prompt_builder=builder)
        .options(num_workers={"CPU": 1})
        .with_column(col("item")),
    )

# ``video_mode="video"`` is vLLM-only, so that one runs on a GPU worker.
ds = ds.with_column(
    "caption_whole",
    VideoCaption(prompt_builder=whole_clip_prompt, video_mode="video")
    .options(num_workers={"GPU": 1})
    .with_column(col("item")),
)

# Sampling and chat-template knobs: one destination each.
ds = ds.with_column(
    "caption_thinking",
    VideoCaption(
        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(1))

Methods

__call__(videos)

Call self as a function.

options(**kwargs)

with_column(videos)

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

Call self as a function.

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