kempnerforge.data.vlm_dataset

VLM dataset and collator (Joint-Decoder).

HuggingFaceVLMDataset wraps a HuggingFace image-text dataset and produces the VLMSample contract:

  • pixel_values: (3, H, W) float tensor, resized to image_size and normalized with the provided mean/std.

  • input_ids: (T,) int64 tensor, right-padded to max_text_len.

  • labels: (T,) int64 tensor matching input_ids with -100 on padding positions and (optionally) on prompt positions when prompt_field is set.

VLMCollator stacks a list of samples into a batch. All batches are padded to the same fixed max_text_len regardless of batch content so different ranks see identical tensor shapes (no NCCL desync under FSDP2). The collator also emits image_positions: (B,) zeros; this slot is reserved for a future multi-image extension and is unused by the Joint-Decoder wrapper today.

The HF datasets and transformers packages are imported lazily so this module is safe to import without them (e.g. in unit tests that don’t exercise the dataset path).

Functions

build_tokenizer(tokenizer_path)

frames_to_clip_tensor(frames, *, max_frames, ...)

Pack frames into a fixed (max_frames, 3, H, W) zero-padded clip + mask.

pil_to_tensor(img, image_size, mean, std)

Resize, convert to (3, H, W) float tensor, and normalize.

resolve_pad_id(tokenizer)

Resolve a pad id, falling back to EOS then 0.

Classes

HuggingFaceVLMDataset

Map-style HF image-text dataset for Joint-Decoder training.

VLMCollator

Stack VLM samples into a fixed-length batch.
kempnerforge.data.vlm_dataset.pil_to_tensor(img, image_size, mean, std)[source]

Resize, convert to (3, H, W) float tensor, and normalize.

Accepts a PIL Image. Converts to RGB if needed so grayscale / RGBA inputs do not drop into the encoder with the wrong channel count.

Public so that out-of-tree callers (e.g. the VLM evaluation adapter) reuse the exact training-time preprocessing as the single source of truth, rather than re-implementing resize/normalize.

Parameters:
Return type:

torch.Tensor

kempnerforge.data.vlm_dataset.frames_to_clip_tensor(frames, *, max_frames, frame_size, image_mean=(0.5, 0.5, 0.5), image_std=(0.5, 0.5, 0.5), dtype=torch.float32)[source]

Pack frames into a fixed (max_frames, 3, H, W) zero-padded clip + mask.

Fills the first min(len(frames), max_frames) slots with per-frame pil_to_tensor output (temporal order preserved) and leaves the remaining slots zero; also returns a (max_frames,) bool frame_mask (True for real frames). A single image is the length-1 case (a 1-frame clip).

Parameters:

  • frames (list[Any]) – Ordered list of PIL Image objects (e.g. decoded clip frames).

  • max_frames (int) – Fixed per-clip frame budget (the output’s frame dimension).

  • frame_size (int) – Square pixel size each frame is resized to.

  • image_std (tuple[float, float, float]) – Per-channel normalization (SigLIP defaults).

  • dtype (torch.dtype) – Output pixel_values dtype.

  • image_mean (tuple[float, float, float])

  • image_std

Returns:

(pixel_values, frame_mask) of shapes (max_frames, 3, frame_size, frame_size) and (max_frames,).

Return type:

tuple[torch.Tensor, torch.Tensor]

kempnerforge.data.vlm_dataset.build_tokenizer(tokenizer_path)[source]
Parameters:

tokenizer_path (str)

Return type:

Any

kempnerforge.data.vlm_dataset.resolve_pad_id(tokenizer)[source]

Resolve a pad id, falling back to EOS then 0.

Mirrors what the dataset uses to right-pad input_ids; tokenizers without a dedicated pad token fall back to the EOS id, and finally to 0 if neither is set.

Parameters:

tokenizer (Any)

Return type:

int

class kempnerforge.data.vlm_dataset.HuggingFaceVLMDataset[source]

Bases: Dataset

Map-style HF image-text dataset for Joint-Decoder training.

Parameters:

  • dataset_name – HF dataset name (e.g. "sayakpaul/coco-30-val-2014") or a local directory written by datasets.save_to_disk.

  • split – Dataset split.

  • image_field – Column name for the PIL image.

  • text_field – Column name for the caption / target text.

  • tokenizer_path – HF tokenizer id or local path.

  • max_text_len – Fixed-length pad target; passed to the collator.

  • prompt_field – Optional column name for a prompt that should NOT receive loss (e.g. the instruction in an instruction-tuned dataset). Prompt tokens get labels=-100.

  • image_size – Target square image size. Default 224.

  • image_std (image_mean /) – Normalization stats. Defaults match SigLIP’s (0.5, 0.5, 0.5).

  • dataset_config – HF dataset config name, if required.

__init__(dataset_name, split, image_field, text_field, tokenizer_path, max_text_len, prompt_field=None, image_size=224, image_mean=(0.5, 0.5, 0.5), image_std=(0.5, 0.5, 0.5), dataset_config=None)[source]
Parameters:
Return type:

None

class kempnerforge.data.vlm_dataset.VLMCollator[source]

Bases: object

Stack VLM samples into a fixed-length batch.

Output keys:

  • pixel_values: (B, 3, H, W).

  • input_ids: (B, max_text_len) int64.

  • labels: (B, max_text_len) int64 with -100 on pad.

  • image_positions: (B,) long tensor. Reserved slot for multi-image extensions; currently all zeros (single image per example placed at sequence position 0).

Padding is always to max_text_len, never batch-max, so ranks always see identical tensor shapes under FSDP2.

__init__(pad_id, max_text_len)[source]
Parameters:

  • pad_id (int)

  • max_text_len (int)

Return type:

None