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; (V, 3, H, W) when tile_grid cuts the image into V views (tile_image).

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

  • labels: (T,) int64 next-token targets (labels[i] = input token i+1), right-padded to max_text_len; pad, trailing, and prompt-predicting positions are -100.

  • text_len: 0-d int64, how many leading positions of input_ids are real tokens (the rest is padding).

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.

PackedVLMCollator packs several documents into each row instead (pack_vlm_samples): a document is its image’s placeholder slots followed by its real text, rows are filled first-fit to seq_len, and the batch carries doc_ids, per-document position_ids and the image placement the Joint-Decoder wrapper uses to write each image’s tokens over its slots.

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)

check_prompt_fits(tokenizer, prompt, ...)

Reject a prompt that leaves no room in max_text_len for a caption token.

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

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

image_to_views(img, image_size, mean, std, *)

One (3, H, W) view, or the (V, 3, H, W) views of tile_image when tiling.

pack_vlm_samples(samples, *, seq_len, ...)

Pack image-text samples into rows of seq_len, each document's image before its text.

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.

tile_image(img, image_size, mean, std, *, ...)

Cut an image into views: the whole image, then an n × n raster of crops.

Classes

HuggingFaceVLMDataset

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

PackedVLMCollator

Pack the documents of a batch into rows of seq_len (pack_vlm_samples).

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.tile_image(img, image_size, mean, std, *, tile_grid, tile_global=True)[source]

Cut an image into views: the whole image, then an n × n raster of crops.

The whole image, resized to image_size square exactly as pil_to_tensor does, comes first when tile_global. The crops are cut from the image resized to tile_grid * image_size square, so each is image_size px and together they cover that canvas exactly; they follow in raster order (left to right, then top to bottom). Returns (views, 3, image_size, image_size), every view normalized like pil_to_tensor.

Public for the same reason as pil_to_tensor: an evaluation adapter must tile a request’s image exactly as training did.

Parameters:
Return type:

torch.Tensor

kempnerforge.data.vlm_dataset.image_to_views(img, image_size, mean, std, *, tile_grid=0, tile_global=True)[source]

One (3, H, W) view, or the (V, 3, H, W) views of tile_image when tiling.

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

kempnerforge.data.vlm_dataset.check_prompt_fits(tokenizer, prompt, max_text_len)[source]

Reject a prompt that leaves no room in max_text_len for a caption token.

Parameters:

  • tokenizer (Any) – The tokenizer the samples are built with.

  • prompt (str) – The prompt, fixed for the whole run.

  • max_text_len (int) – The fixed text length of every sample.

Raises:

ValueError – max_text_len is below the prompt’s token count plus one caption token, plus one more when the tokenizer defines an EOS.

Return type:

None

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 be supervised (e.g. the instruction in an instruction-tuned dataset). Positions predicting prompt tokens are -100; the last prompt token still predicts the first caption token. A row whose prompt leaves no room in max_text_len for a caption token has every label masked, and building the dataset raises ValueError when no row leaves that room. A None value means the row has no prompt. Building the dataset raises TypeError unless the column holds strings (dtype string or large_string, or null when no row has a prompt).

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

  • tile_grid – 0 keeps one view per image; n cuts each image into an n × n grid of image_size crops (tile_image), so pixel_values is (views, 3, H, W).

  • tile_global – Whether a tiled image’s views start with the whole image.

__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, tile_grid=0, tile_global=True)[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), or (B, V, 3, H, W) when every sample carries V views of its image.

  • 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

kempnerforge.data.vlm_dataset.pack_vlm_samples(samples, *, seq_len, tokens_per_image, pad_id)[source]

Pack image-text samples into rows of seq_len, each document’s image before its text.

A sample is one document: pixel_values (3, H, W) or (V, 3, H, W) (V views, each an encoder input), input_ids and labels (T,) whose first text_len positions are real (0 packs the image alone with nothing to supervise, as the unpacked path’s all-masked row does). Its footprint is V * tokens_per_image placeholder slots followed by text_len tokens. The last packed label is forced to -100 as a guard for samples built outside HuggingFaceVLMDataset (which already leaves it so): the token after a trimmed text_len is not in the row. Documents go first-fit into rows in order: into the first open row with room, else a new row. A row’s tail is padding that forms its own document.

Returns:

pixel_values (N, 3, H, W), every image or view across the batch in document order; input_ids, labels and doc_ids (rows, seq_len), with pad_id and -100 on the image slots and the padding tail, and doc_ids numbering a row’s documents from 0 with the tail one past the last; position_ids (rows, seq_len), each document’s positions (and the tail’s) restarting at 0; image_rows and image_starts (N,), the row and first slot of each image; image_tokens, 0-d, the tokens_per_image the slots were laid out with; text_rows and text_positions (M,), the row and position of every real text token in document order, where the wrapper runs the LM head.

Raises:

ValueError – samples is empty, a sample lacks text_len, a document has no image or a footprint over seq_len, or seq_len or tokens_per_image is not positive.

Parameters:
Return type:

dict[str, torch.Tensor]

class kempnerforge.data.vlm_dataset.PackedVLMCollator[source]

Bases: object

Pack the documents of a batch into rows of seq_len (pack_vlm_samples).

The number of rows follows the documents’ footprints, so it varies from batch to batch; train.batch_size is the number of documents per micro-batch, not the number of rows.

__init__(pad_id, seq_len, tokens_per_image)[source]
Parameters:

  • pad_id (int)

  • seq_len (int)

  • tokens_per_image (int)

Return type:

None