diff --git a/CHANGELOG.md b/CHANGELOG.md index 49ab032..6ab0c87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added + - **MoE router z-loss** (`moe_router_z_loss_weight`, ST-MoE style). An optional penalty on the router's pre-softmax logits — per MoE layer `mean_token(logsumexp(router_logits))²` — summed across layers and added to the training loss as `moe_router_z_loss_weight × z_loss`. It keeps router logits from growing without bound, targeting *logit-growth stability* (not load balance — that's the aux loss). Default `0.0` is off: the term is never added, so training, outputs, and gradients are unchanged. `z_loss` is a plain attribute like `aux_loss` (not a buffer/parameter), so it never enters `state_dict` — checkpoint-safe. - `kempnerforge/config/model.py`: `moe_router_z_loss_weight: float = 0.0` (with a non-negativity check). - `kempnerforge/model/router.py`: both `SoftmaxTopKRouter` and `SigmoidTopKRouter` set `self.z_loss = (logsumexp(logits, dim=-1) ** 2).mean()`. @@ -78,14 +79,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Milestone-aware retention: `CheckpointManager._cleanup` never prunes a step where the configured dynamic strategy fired, so `keep_last_n` rotates only the later interval checkpoints. `keep_last_n <= 0` keeps everything (the previous `keep_last_n >= 1` requirement is relaxed). - `scripts/train.py`: the save gate now calls `config.checkpoint.should_save(step)`. - Tests: `tests/unit/test_config.py` (defaults, power2 firing, offset-based `start > 0`, validation, unknown-strategy rejection, `is_dynamic_milestone`), `tests/unit/test_checkpoint.py::TestCheckpointRetention::test_cleanup_protects_dynamic_milestones`. +- **VLM evaluation pipeline** (`scripts/vlm_eval_harness.py` + `kempnerforge/eval/vlm/`). Evaluates any KempnerForge VLM checkpoint on any standard multimodal benchmarks (MMMU, MMBench, ScienceQA, SEED, AI2D, …) by integrating the [lmms-eval](https://github.com/EvolvingLMMs-Lab/lmms-eval) harness through a custom model adapter that wraps `VLMWrapper` and loads directly from a DCP checkpoint. Arch-agnostic across Joint-Decoder / Cross-Attention / MoT; MoMa fails fast (its non-causal expert-choice routing cannot autoregressively generate, and eval requires generation). v1 is single-GPU, image-only, and generation-only (`generate_until`), decoding requests in batches (the `--batch-size` model-arg). All changes are additive and backward compatible; the only edit to existing code is a behavior-preserving refactor. + - `kempnerforge/eval/__init__.py`, `kempnerforge/eval/vlm/__init__.py`: new eval-subsystem namespace. Import-isolated — neither is imported on the default `import kempnerforge` path, so the main package keeps working with lmms-eval absent (pinned by `test_import_isolation.py`). + - `kempnerforge/eval/vlm/adapter.py`: `KempnerForgeVLM(lmms)` chat adapter (`is_simple = False`). Loader (`build_vlm_wrapper` behind a `_build_model` seam → single-process `dcp.load` of model shards only; `resolve_resume_path` with a specific-`step_N` fallback; reads plain-JSON `metadata.json`, never `train_state.pt`); prompt rendering by flattening `ChatMessages` text blocks (no chat template, no `` placeholder — images are conditioned via `pixel_values`); a cache-less, **batched** greedy/sampled decode loop (`_generate_batch`, reusing `kempnerforge.model.generate.sample`): `generate_until` groups requests by `gen_kwargs` and chunks/reorders them via `lmms_eval.utils.Collator` (mirrors `lmms_eval/models/chat/qwen2_5_vl.py`), **right-padding** the text to the batch-max length (image prefix at `0..n-1`, text contiguous from `n`, trailing pads causally masked — the training layout; reuses the now-public `resolve_pad_id` for pad ids) so a batched forward gives each row the same real-position logits as decoding it alone (pinned by a batch-equivalence test), reading each row's next token at its own last real position and tracking EOS / `until` / `max_new_tokens` per row. `generate_until` only — `loglikelihood` and `generate_until_multi_round` raise `NotImplementedError`. Guards (clear `NotImplementedError`/`ValueError`): MoMa arch, video/audio, multi-image, multi-turn/few-shot. A file-level `# pyright: reportMissingImports=false` keeps `pyright kempnerforge/` green in CI (where the undeclared lmms-eval is absent) without an inline ignore. + - `kempnerforge/eval/vlm/manifest.py`: `MANIFEST = ModelManifest(model_id="kempnerforge_vlm", chat_class_path="kempnerforge.eval.vlm.adapter.KempnerForgeVLM")` for lmms-eval entry-point discovery. + - `scripts/vlm_eval_harness.py`: CLI mirroring `scripts/eval_harness.py` (no conversion). `--config`/`--checkpoint`/`--tasks` required (default suite TBD), plus `--limit`/`--output`/`--device`/`--dtype`/`--batch-size`/`--max-new-tokens`; lazy `lmms_eval.evaluator.simple_evaluate` import with a helpful error. + - `pyproject.toml`: `[project.entry-points."lmms_eval.models"]` for the adapter — metadata only; lmms-eval is NOT added as a dependency (install separately with `uv pip install lmms-eval`, mirroring how lm-eval is handled). + - `kempnerforge/data/vlm_dataset.py`: behavior-preserving refactor — `_pil_to_tensor` → public `pil_to_tensor`; tokenizer construction and pad-id resolution extracted to public `build_tokenizer` / `resolve_pad_id`, so the eval adapter reuses the exact training-time preprocessing as the single source of truth. (`tests/unit/test_vlm_dataset.py` updated to the renamed helper; all behavior identical.) + - Tests: `tests/unit/eval/vlm/` (CPU) runs against a faithful in-repo fake `lmms_eval` injected via `conftest.py` (`_fake_lmms_eval.py`) — so the adapter/registry tests **always run in CI and contribute coverage** without the undeclared lmms-eval dependency, instead of skipping. `test_adapter.py` covers rendering / preprocessing / `gen_kwargs` / the batched decode loop (`TestGenerateBatchSingle` B=1 regressions + `TestGenerateBatchMulti` batch-equivalence, per-row `max_new_tokens`, per-row EOS) / guards, plus the loader helpers (`_resolve_dtype`, `_load_config`, `_load_weights`, `_log_checkpoint_metadata`, `_first_stop`), the `__init__` guards, and `generate_until` end-to-end; `test_manigest.py` covers the manifest. `test_import_isolation.py` still asserts `import kempnerforge` needs no lmms-eval. `tests/integration/test_lmms_eval_contract.py` pins the *real* lmms-eval API and entry-point resolution to the fakes' assumptions (gated on real lmms-eval), and `tests/integration/test_vlm_eval.py` keeps the self-contained DCP round-trip (at `batch_size=2`) + env-gated real-task path. `tests/conftest.py`: `tiny_vlm_configs` / `tiny_vlm_wrapper` fixtures (random vision encoder, CPU, no checkpoint). + - Docs: `docs/how-to/run-vlm-evaluation.md`, wired into the how-to `toctree`. + - Deferred: data-parallel and sharded multi-GPU inference, a representative default benchmark suite, whether to formalize the lmms-eval dependency, and confirming frozen vision-encoder weights load from a real checkpoint. ### Changed + - `docs/getting-started/install.md` Prerequisites: documents `.python-version` and uv's auto-fetch behavior. - `README.md` and `kempnerforge/README.md` Prerequisites: clarify that uv auto-fetches Python 3.12 via `.python-version`. - `docs/claude-ready.md` first-run flow: `/kempnerforge:install-and-verify` runs before `/kempnerforge:cluster-config`. - `README.md` and `kempnerforge/README.md`: list `install-and-verify` in the skill catalog and drop the hardcoded skill count. ### Fixed + - **Resume silently reset AdamW optimizer momentum.** `CheckpointManager` round-tripped optimizer state through raw `optimizer.state_dict()` / `optimizer.load_state_dict()`. On resume the optimizer is freshly built, so its `state_dict()` is empty — `dcp.load` then had no `exp_avg` / `exp_avg_sq` tensors to fill, and the moments were silently dropped, resetting Adam momentum to zero at every resume point. Model weights, scheduler, dataloader position, and RNG all restored correctly; only the optimizer moments were lost, so resumed runs were not bit-exact. - `kempnerforge/checkpoint/manager.py`: save and load now go through DCP's `get_model_state_dict` / `get_optimizer_state_dict` / `set_model_state_dict` / `set_optimizer_state_dict`. The getters build a load template with the optimizer moments allocated in the correct FSDP/DTensor layout, so `dcp.load` repopulates them; the setters write the loaded values back into the live optimizer. - `docs/checkpointing/dcp-model.md`: updated the save/load snippets and the "shape to fill" explanation to the DCP-aware helpers. @@ -97,6 +110,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Initial public release. ### Architecture + - Decoder-only Transformer with RoPE, GQA, SwiGLU MLP, RMSNorm - Optional QK-Norm (Gemma/DeepSeek-V3 style) - Mixture-of-Experts (MoE) with softmax top-k and DeepSeek-V3 sigmoid routers @@ -105,6 +119,7 @@ Initial public release. - Sequence-level aux loss, gradient scaling, adaptive bias schedule for sigmoid router ### Parallelism + - **FSDP2** — composable `fully_shard()`, per-block sharding, mixed precision (bf16/fp16/fp32) - **Tensor Parallelism** — column/row parallel with SequenceParallel and meta-device init - **Expert Parallelism** — all-to-all dispatch, multi-node EP+TP+FSDP2 composition @@ -113,6 +128,7 @@ Initial public release. - **SDPA backend override** — `model.sdpa_backend` config (`auto`/`flash`/`efficient`/`cudnn`/`math`) for kernel benchmarking and debugging ### Training + - Optimizers: AdamW, Muon (Newton-Schulz orthogonalized momentum), Lion (half optimizer memory), Schedule-Free AdamW - LR schedulers: cosine, linear, WSD (cosine/linear/sqrt cooldown), constant, REX, none - Loss functions: cross-entropy, chunked cross-entropy, z-loss regularizer @@ -124,29 +140,34 @@ Initial public release. - Training loop hooks (`TrainingHook`, `HookRunner`) for extensibility ### Resilience + - SLURM preemption handling (SIGTERM/SIGUSR1) with cooperative shutdown - NaN detection with configurable actions (warn, skip, raise) - GPU compute/memory health checks - NCCL liveness monitoring via all-reduce ### Interpretability + - Activation extraction hooks (`ActivationStore`) with CPU offload - Attention weight capture (explicit QK^T path for mechanistic interpretability) - Batch extraction over datasets for CKA/SVCCA/probing pipelines ### Metrics & Observability + - Per-step MetricsTracker with EMA smoothing - MFU computation from architecture params and GPU peak FLOPs - WandB and TensorBoard backends - Peak memory monitoring with optional snapshot export ### Configuration + - Typed dataclass configs per domain (`ModelConfig`, `TrainConfig`, etc.) - Layered loading: defaults → TOML file → CLI overrides - Fail-fast validation at object construction - Registry for swappable components (optimizers, schedulers, losses, routers, norms) ### Testing + - 794 unit tests (CPU-only) - Integration tests (1 GPU) - Distributed tests (4 GPUs via torchrun) diff --git a/docs/how-to/index.md b/docs/how-to/index.md index d44ff02..2e914c4 100644 --- a/docs/how-to/index.md +++ b/docs/how-to/index.md @@ -28,6 +28,8 @@ with runnable code (or a link to a notebook, config, or script that runs it) - [Run evaluation](run-evaluation.md) — `run_eval()` during training, standalone `scripts/eval.py`, `scripts/eval_harness.py` for lm-eval-harness. +- [Run VLM evaluation](run-vlm-evaluation.md) — multimodal benchmarks via + lmms-eval, on any VLM DCP checkpoint. - [Generate from checkpoint](generate-from-checkpoint.md) — load a DCP checkpoint, call `generate()` with temperature / top-k / top-p, interact with `KVCache`. @@ -68,6 +70,7 @@ end-to-end-training-run scaling-guide slurm-distributed-setup run-evaluation +run-vlm-evaluation generate-from-checkpoint debug-training-regressions compare-optimizers diff --git a/docs/how-to/run-vlm-evaluation.md b/docs/how-to/run-vlm-evaluation.md new file mode 100644 index 0000000..4e61d59 --- /dev/null +++ b/docs/how-to/run-vlm-evaluation.md @@ -0,0 +1,184 @@ +# Run VLM evaluation + +Evaluate a vision-language model (VLM) checkpoint on any standard multimodal +benchmark (MMMU, MMBench, ScienceQA, SEED, AI2D, …) by integrating the +[lmms-eval](https://github.com/EvolvingLMMs-Lab/lmms-eval) harness. + +[Run evaluation](run-evaluation.md). + +| Entry point | When to use | +|-------------|-------------| +| `scripts/vlm_eval_harness.py` | Downstream VLM benchmarks via lmms-eval, on any DCP checkpoint | + +A custom lmms-eval *chat model* loads `VLMWrapper` directly from the DCP checkpoint. The pieces are: +[`kempnerforge/eval/vlm/adapter.py`](https://github.com/KempnerInstitute/KempnerForge/blob/main/kempnerforge/eval/vlm/adapter.py) +(the `KempnerForgeVLM` adapter), +[`kempnerforge/eval/vlm/manifest.py`](https://github.com/KempnerInstitute/KempnerForge/blob/main/kempnerforge/eval/vlm/manifest.py) +(the lmms-eval registration manifest), and +[`scripts/vlm_eval_harness.py`](https://github.com/KempnerInstitute/KempnerForge/blob/main/scripts/vlm_eval_harness.py) +(the CLI). + +## What's on this page + +- [Install lmms-eval](#install-lmms-eval) — the optional, separately-installed dependency +- [Usage](#usage) — running the harness on a checkpoint +- [Flags](#flags) — CLI options at a glance +- [Video evaluation](#video-evaluation) — evaluating a video checkpoint +- [Limitations](#limitations) — single-GPU, MoMa, multi-turn/few-shot/multi-image, flattening, no KV cache +- [Cluster environment notes](#cluster-environment-notes) — torchvision / libstdc++ gotchas + +## Install lmms-eval + +`lmms-eval` is an **optional dependency** and is intentionally NOT declared in +`pyproject.toml`. Install +it into your environment before running: + +```bash +uv pip install lmms-eval +``` + +The `lmms_eval.models` entry point that exposes the `kempnerforge_vlm` model is +declared in `pyproject.toml` as metadata only — it does not pull lmms-eval in as +a dependency, and `import kempnerforge` works without lmms-eval installed. + +**Video evaluation** additionally needs the `av` (PyAV) video-decoding +dependency, which ships in the optional `video` group: + +```bash +uv sync --group video +``` + +PyAV's manylinux wheel bundles FFmpeg, so no system FFmpeg or CUDA libraries are +required. (Image-only evaluation does not need this group.) + +## Usage + +```bash +# One task, write results JSON +uv run python scripts/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val \ + --output results/vlm_step_10000.json + +# Several tasks, quick partial run (4 examples per task) +uv run python scripts/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val,mmbench_en_dev,scienceqa_img \ + --limit 4 +``` + +`--config` is the same KempnerForge TOML the checkpoint was trained with (it +carries the vision encoder, adapter, `vlm.arch`, and tokenizer settings). +`--checkpoint` accepts either a run directory (the latest `step_N` is resolved +automatically) or a specific `step_N` directory. + +There is **no default task suite** — `--tasks` is required. A representative +default benchmark set is still being decided. + +## Flags + +| Flag | Default | Purpose | +|------|---------|---------| +| `--config` | — (required) | KempnerForge TOML the checkpoint was trained with | +| `--checkpoint` | — (required) | DCP checkpoint dir (run dir or `step_N` dir) | +| `--tasks` | — (required) | comma-separated lmms-eval task names | +| `--limit` | `None` | cap examples per task (int count, or `<1.0` fraction) | +| `--output` | `None` | save full JSON results | +| `--device` | `cuda` | inference device | +| `--dtype` | `None`(maps to model config setting) | model dtype | +| `--batch-size` | `1` | requests decoded together (grouped by `gen_kwargs`) | +| `--max-new-tokens` | `128` | fallback only; task `gen_kwargs` override it | + +## Video evaluation + +When `--config` is a **video checkpoint** (its TOML has a `[video]` section), the +harness evaluates lmms-eval *video* `generate_until` tasks: each request's video +is decoded into frames and fed to the model as a single clip. This needs the `av` +video group (see [Install lmms-eval](#install-lmms-eval)). + +```bash +uv run python scripts/vlm_eval_harness.py \ + --config configs/train/vlm_video_webvid.toml \ + --checkpoint checkpoints/vlm_video/step_10000 \ + --tasks \ + --limit 4 +``` + +- **The frame budget is a property of the checkpoint, not a flag.** Frames are + sampled by the model's own `[video]` policy (`fps` / `min_frames` / + `max_frames`, the Molmo2 uniform `sample_timestamps`) and fixed to exactly + `max_frames` (zero-padded when a clip yields fewer). You cannot change it at + eval time — the transformer was built around `frames_per_clip = max_frames`. + Comparability to externally published video-benchmark numbers therefore depends + on the checkpoint's frame budget matching the reference's, which is a training + choice rather than a knob here. +- **Scope.** One video per request, single-turn, zero-shot, generative arches + (`joint_decoder` / `cross_attention` / `mot`). A single **image** task also runs + on a video checkpoint — the image is treated as a 1-frame clip, zero-padded to + `frames_per_clip`. Multiple videos, mixed image+video, multiple images, audio, + and multi-turn / few-shot raise a clear error; MoMa still fails fast. An + **image** checkpoint cannot evaluate video and raises a clear error if handed a + video task. + +## Limitations + +Several are tracked follow-ups. + +- **Single GPU.** v1 runs on one GPU. Data-parallel + multi-GPU is a localized + future addition; sharded/model-parallel inference for models too large for one + GPU is a larger, separate effort. +- **MoMa is not supported.** The `moma` arch uses non-causal expert-choice + routing and cannot autoregressively generate, but eval tasks are + generation-only. A MoMa checkpoint fails fast with a clear error. Joint-Decoder + (`joint_decoder`), Cross-Attention (`cross_attention`), and MoT (`mot`) are + supported. +- **One visual per request; no multi-turn / few-shot / multi-image.** A request + carries exactly one image (image checkpoint) or one video (video checkpoint — + see [Video evaluation](#video-evaluation)). Audio, multiple images, multiple + videos, mixed image+video, and multi-turn / few-shot requests raise a clear + error. Multi-image and multi-turn/few-shot are tracked follow-ups (for chat + tasks lmms-eval delivers few-shot as extra content blocks/turns, so it reduces + to multi-image + multi-turn support). +- **Prompt flattening discards structure.** Flattening drops role/turn structure + and any model-specific chat template. KempnerForge pre-training uses no chat + template; once a post-training format exists, repo-wide chat-template support + should be added and the rendering step made configurable. +- **No KV cache.** Decoding re-runs the full transformer over the growing sequence + each step (KempnerForge has no image-conditioned KV-cache decode path); this is + correct but costs extra compute, and a KV-cache decode is future work. Raising + `--batch-size` decodes multiple requests together + (right-padded, grouped by `gen_kwargs`) to amortize the per-step transformer cost. + +## Cluster environment notes + +Installing lmms-eval pulls in extra packages that can clash with a CUDA-pinned +PyTorch. Two gotchas seen on the Kempner cluster: + +- **torchvision must match the CUDA build of torch.** The default-index + `torchvision` is ABI-incompatible with `torch …+cu128` (it fails + `register_fake("torchvision::nms")`, which breaks `import lmms_eval`). Install + the matching build from the same index: + + ```bash + uv pip install --reinstall-package torchvision \ + --index https://download.pytorch.org/whl/cu128 "torchvision==0.26.0" + ``` + +- **`GLIBCXX_… not found` when importing the evaluator.** lmms-eval's + `simple_evaluate` pulls in a library that needs a newer `libstdc++` than the + system one. Put a newer `libstdc++` first on the library path, e.g. + `LD_LIBRARY_PATH=/lib uv run python scripts/vlm_eval_harness.py …`. + +## See also + +- [Run evaluation](run-evaluation.md) — text-model loss/perplexity and the + `lm-eval` harness this page parallels. +- [End-to-end training run](end-to-end-training-run.md) — produces the + checkpoints this harness consumes. +- [`kempnerforge/eval/vlm/adapter.py`](https://github.com/KempnerInstitute/KempnerForge/blob/main/kempnerforge/eval/vlm/adapter.py) + and + [`scripts/vlm_eval_harness.py`](https://github.com/KempnerInstitute/KempnerForge/blob/main/scripts/vlm_eval_harness.py) + — the adapter and CLI this page documents. diff --git a/kempnerforge/config/vlm.py b/kempnerforge/config/vlm.py index 777c27f..f3640c2 100644 --- a/kempnerforge/config/vlm.py +++ b/kempnerforge/config/vlm.py @@ -156,6 +156,18 @@ def residual_stream_image_tokens(self, num_tokens: int) -> int: """ return num_tokens + @property + def is_generative(self) -> bool: + """Whether this arch can autoregressively generate token-by-token. + + Generation-only consumers (e.g. the lmms-eval chat adapter in + ``kempnerforge/eval/vlm``) query this to fail fast on arches that + cannot decode autoregressively. Defaults to ``True`` (the common + case); a non-causal arch overrides it to ``False`` (see + ``MoMaConfig``). + """ + return True + @classmethod def for_arch(cls, arch: str, **kwargs: Any) -> VLMConfig: """Resolve ``arch`` to a registered subclass and instantiate. @@ -466,6 +478,12 @@ def residual_stream_image_tokens(self, num_tokens: int) -> int: """ return num_tokens + @property + def is_generative(self) -> bool: + # Expert-choice routing is non-causal (see the class docstring), so MoMa + # cannot autoregressively generate; generation-only consumers reject it. + return False + def effective_capacity_factor(self, modality: str) -> float: """Resolve the per-expert capacity factor for ``modality``. diff --git a/kempnerforge/data/video_dataset.py b/kempnerforge/data/video_dataset.py index ff459cc..b84281e 100644 --- a/kempnerforge/data/video_dataset.py +++ b/kempnerforge/data/video_dataset.py @@ -36,8 +36,9 @@ from kempnerforge.data.vlm_dataset import ( DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD, - _pil_to_tensor, _tokenize_and_mask, + frames_to_clip_tensor, + resolve_pad_id, ) logger = logging.getLogger(__name__) @@ -48,13 +49,6 @@ _VIDEO_SUBDIR = {"train": "train", "validation": "validation"} -def _resolve_pad_id(tokenizer: Any) -> int: - pad_id = tokenizer.pad_token_id - if pad_id is None: - pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0 - return int(pad_id) - - class VideoDataset(Dataset): """Base for video-caption datasets feeding the VLM video path. @@ -122,7 +116,7 @@ def __init__( ) self._ids, self._caps = self._load_manifest(csv_dir, max_samples) self._tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - self._pad_id = _resolve_pad_id(self._tokenizer) + self._pad_id = resolve_pad_id(self._tokenizer) self._max_text_len = max_text_len self._max_frames = max_frames self._min_frames = min_frames @@ -202,18 +196,17 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: logger.debug("video decode failed for %s: %s", path, e) frames = [] - f = self._max_frames - size = self._frame_size - pixel_values = torch.zeros(f, 3, size, size, dtype=torch.float32) - frame_mask = torch.zeros(f, dtype=torch.bool) - n_real = min(len(frames), f) - for i in range(n_real): - pixel_values[i] = _pil_to_tensor(frames[i], size, self._image_mean, self._image_std) - frame_mask[i] = True + pixel_values, frame_mask = frames_to_clip_tensor( + frames, + max_frames=self._max_frames, + frame_size=self._frame_size, + image_mean=self._image_mean, + image_std=self._image_std, + ) prompt = self._prompt or None input_ids, labels = _tokenize_and_mask(self._tokenizer, caption, self._max_text_len, prompt) - if n_real == 0: + if not frames: # Undecodable clip: keep static shapes but contribute no loss. labels = torch.full_like(labels, -100) return { diff --git a/kempnerforge/data/video_io.py b/kempnerforge/data/video_io.py index b6d91ab..056941a 100644 --- a/kempnerforge/data/video_io.py +++ b/kempnerforge/data/video_io.py @@ -17,7 +17,7 @@ decoding requires the package. Returned frames are ``PIL.Image`` objects so the caller can reuse the exact -image preprocessing (``_pil_to_tensor``) used on the single-image path. +image preprocessing (``pil_to_tensor``) used on the single-image path. """ from __future__ import annotations diff --git a/kempnerforge/data/vlm_dataset.py b/kempnerforge/data/vlm_dataset.py index ec6059e..4d5a746 100644 --- a/kempnerforge/data/vlm_dataset.py +++ b/kempnerforge/data/vlm_dataset.py @@ -38,7 +38,7 @@ DEFAULT_IMAGE_STD = (0.5, 0.5, 0.5) -def _pil_to_tensor( +def pil_to_tensor( img: Any, image_size: int, mean: tuple[float, float, float], @@ -49,6 +49,10 @@ def _pil_to_tensor( 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. """ from PIL import Image @@ -67,6 +71,61 @@ def _pil_to_tensor( return (t - mean_t) / std_t +def frames_to_clip_tensor( + frames: list[Any], + *, + max_frames: int, + frame_size: int, + image_mean: tuple[float, float, float] = DEFAULT_IMAGE_MEAN, + image_std: tuple[float, float, float] = DEFAULT_IMAGE_STD, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """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). + + Args: + frames: Ordered list of PIL ``Image`` objects (e.g. decoded clip frames). + max_frames: Fixed per-clip frame budget (the output's frame dimension). + frame_size: Square pixel size each frame is resized to. + image_mean / image_std: Per-channel normalization (SigLIP defaults). + dtype: Output ``pixel_values`` dtype. + + Returns: + ``(pixel_values, frame_mask)`` of shapes ``(max_frames, 3, frame_size, + frame_size)`` and ``(max_frames,)``. + """ + pixel_values = torch.zeros(max_frames, 3, frame_size, frame_size, dtype=dtype) + frame_mask = torch.zeros(max_frames, dtype=torch.bool) + n_real = min(len(frames), max_frames) + for i in range(n_real): + pixel_values[i] = pil_to_tensor(frames[i], frame_size, image_mean, image_std) + frame_mask[i] = True + return pixel_values, frame_mask + + +def build_tokenizer(tokenizer_path: str) -> Any: + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(tokenizer_path) + + +def resolve_pad_id(tokenizer: Any) -> int: + """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. + """ + pad_id = tokenizer.pad_token_id + if pad_id is None: + pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0 + return pad_id + + def _tokenize_and_mask( tokenizer: Any, text: str, @@ -99,9 +158,7 @@ def _tokenize_and_mask( prompt_len = 0 full_ids = full_ids[:max_text_len] - pad_id = tokenizer.pad_token_id - if pad_id is None: - pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0 + pad_id = resolve_pad_id(tokenizer) n = len(full_ids) input_ids = torch.full((max_text_len,), pad_id, dtype=torch.long) @@ -153,7 +210,6 @@ def __init__( from datasets import Dataset as HFDataset from datasets import load_dataset, load_from_disk - from transformers import AutoTokenizer # When dataset_name points at an existing directory on disk (absolute # or relative path), prefer load_from_disk. Otherwise treat it as an @@ -183,7 +239,7 @@ def __init__( self._image_mean = image_mean self._image_std = image_std self._max_text_len = max_text_len - self._tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + self._tokenizer = build_tokenizer(tokenizer_path) logger.info( f"HuggingFaceVLMDataset: {dataset_name} [{split}], {len(self._ds):,} samples, " f"image_size={image_size}, max_text_len={max_text_len}" @@ -194,7 +250,7 @@ def __len__(self) -> int: def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: row = self._ds[idx] - pixels = _pil_to_tensor( + pixels = pil_to_tensor( row[self._image_field], self._image_size, self._image_mean, self._image_std ) prompt_val = row[self._prompt_field] if self._prompt_field is not None else None diff --git a/kempnerforge/eval/__init__.py b/kempnerforge/eval/__init__.py new file mode 100644 index 0000000..17fb5de --- /dev/null +++ b/kempnerforge/eval/__init__.py @@ -0,0 +1,13 @@ +"""Evaluation subsystems for KempnerForge. + +This namespace hosts model-type-specific evaluation subsystems that integrate +external harnesses. ``eval/vlm/`` (VLM evaluation via lmms-eval) is the only +subsystem today; future siblings (``eval/lm/``, ``eval/audio/``, ...) slot in +without restructuring. + +Import isolation: this package and its submodules are deliberately NOT imported +from ``kempnerforge/__init__.py`` or any default import path, because +``eval/vlm`` depends on the optional, undeclared ``lmms-eval`` package. +``import kempnerforge`` must keep working with ``lmms-eval`` absent, so do not +add eager imports of the eval subpackages here. +""" diff --git a/kempnerforge/eval/vlm/__init__.py b/kempnerforge/eval/vlm/__init__.py new file mode 100644 index 0000000..f758ab1 --- /dev/null +++ b/kempnerforge/eval/vlm/__init__.py @@ -0,0 +1,9 @@ +"""VLM evaluation subsystem: an lmms-eval chat-model adapter over ``VLMWrapper``. + +The adapter (``adapter.py``) and its registration manifest (``manifest.py``) +both import the optional ``lmms-eval`` package, so they are intentionally NOT +imported here. They are loaded only by ``scripts/vlm_eval_harness.py`` and by +lmms-eval's entry-point loader (at which point ``lmms-eval`` is necessarily +installed). Keeping this ``__init__`` import-free preserves +``import kempnerforge`` with ``lmms-eval`` absent. +""" diff --git a/kempnerforge/eval/vlm/adapter.py b/kempnerforge/eval/vlm/adapter.py new file mode 100644 index 0000000..b4228d6 --- /dev/null +++ b/kempnerforge/eval/vlm/adapter.py @@ -0,0 +1,705 @@ +# pyright: reportMissingImports=false +# ^ lmms-eval is an optional, UNDECLARED dependency (it is installed separately, +# not listed in pyproject.toml). CI type-checks `kempnerforge/` without it +# installed, so the `lmms_eval` imports below would otherwise raise +# reportMissingImports. A file-level directive (not a `# type: ignore`, which +# reportUnnecessaryTypeIgnoreComment would flag as unnecessary in dev where +# lmms-eval *is* installed) scopes the relaxation to this one module. +"""lmms-eval chat-model adapter wrapping a KempnerForge ``VLMWrapper``. + +This module implements ``KempnerForgeVLM``, an lmms-eval ``chat`` model +(``is_simple = False``) that evaluates any KempnerForge VLM checkpoint on the +standard multimodal benchmarks lmms-eval implements as ``generate_until`` tasks +(MMMU, MMBench, ScienceQA, SEED, AI2D, ...). It is loaded directly from a DCP checkpoint, +and is arch-agnostic across the generative VLM arches. +v1 scope and deliberate choices (see docs/how-to/run-vlm-evaluation.md): + +- **Generation: no transformer KV cache, single-GPU, batched.** The decode loop + re-runs the transformer (including the vision encoder + adapter) over the + growing sequence each step. There is no transformer KV cache + (``Transformer.forward`` forbids combining ``kv_caches`` with any + image-conditioning route), and KempnerForge has no image-conditioned KV-cache + decode path. Requests are decoded in batches + (``batch_size`` model-arg) by **right-padding** the text to the batch-max + length — the same layout training uses (image prefix at ``0..n-1``, text + contiguous from ``n``, trailing pads causally masked) — and reading each + row's logits at its own last real position. Single-GPU is the validated + invocation, not a baked-in assumption: rank/world_size come from the lmms + base (defaults 0/1) and model construction sits behind ``_build_model`` so a + data-parallel path is a localized future change. + +- **Prompt rendering: flatten, no chat template.** KempnerForge pre-training + uses no chat template / processor and no ```` placeholder (images are + conditioned at the embedding level). We render an lmms-eval ``ChatMessages`` + by concatenating its text content blocks in order into a single prompt + string. This discards role/turn structure and any model-specific template. + A future enhancement should add repo-wide chat-template support (applied once + a post-training format exists), at which point this rendering step becomes + configurable rather than hard-coded to flattening. + +- **Arch coverage.** Joint-Decoder, Cross-Attention, and MoT are supported. + MoMa is NOT: its expert-choice routing is non-causal and cannot + autoregressively generate, and chat tasks are generation-only. A MoMa + checkpoint fails fast in ``__init__``. + +- **Image and video.** An image checkpoint evaluates exactly one image per + request; a video checkpoint (a ``[video]`` config) evaluates one video per + request — decoded to a fixed ``frames_per_clip`` clip via the training + frame-sampling policy — and also accepts a single image (a 1-frame clip). + Audio, multi-image, multiple videos, mixed image+video, and multi-turn/few-shot + requests raise ``NotImplementedError``; ``loglikelihood`` and + ``generate_until_multi_round`` are not implemented (chat tasks are + generation-only). Visual input is modeled as an ordered list of frames (a + single image is the length-1 case). +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import torch +import torch.distributed.checkpoint as dcp +from lmms_eval.api.instance import Instance +from lmms_eval.api.model import lmms +from lmms_eval.protocol import ChatMessages +from lmms_eval.utils import Collator +from tqdm import tqdm + +from kempnerforge.config.job import JobConfig +from kempnerforge.config.loader import load_config +from kempnerforge.config.video import VideoConfig +from kempnerforge.config.vlm import VLMConfig +from kempnerforge.data.video_io import decode_video_frames +from kempnerforge.data.vlm_dataset import ( + DEFAULT_IMAGE_MEAN, + DEFAULT_IMAGE_STD, + build_tokenizer, + frames_to_clip_tensor, + pil_to_tensor, + resolve_pad_id, +) +from kempnerforge.metrics.logger import get_logger +from kempnerforge.model.generate import sample +from kempnerforge.model.vlm import VLMWrapper, build_vlm_wrapper +from kempnerforge.resilience.elastic import resolve_resume_path + +logger = get_logger(__name__) + +DEFAULT_MAX_NEW_TOKENS = 128 + +_DTYPES = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, +} + + +def _resolve_dtype(dtype: str | torch.dtype) -> torch.dtype: + if isinstance(dtype, torch.dtype): + return dtype + try: + return _DTYPES[dtype] + except KeyError: + raise ValueError(f"Unsupported dtype {dtype!r}; choose from {sorted(_DTYPES)}") from None + + +# --------------------------------------------------------------------------- # +# Loader +# --------------------------------------------------------------------------- # + + +def _build_model(config: JobConfig, device: torch.device, dtype: torch.dtype) -> VLMWrapper: + assert config.vlm is not None, "internal: _build_model requires a VLM config" + assert config.vision_encoder is not None, "internal: VLM config requires a vision encoder" + assert config.adapter is not None, "internal: VLM config materializes a default adapter" + # A video checkpoint bakes num_image_tokens = frames_per_clip * tokens_per_frame + # into the transformer's residual/positional structure, so the eval must rebuild + # with the same frames_per_clip the training run used (mirrors scripts/train.py). + # Image configs default to 1. + frames_per_clip = config.video.max_frames if config.video is not None else 1 + model = build_vlm_wrapper( + config.model, + config.vision_encoder, + config.adapter, + config.vlm, + frames_per_clip=frames_per_clip, + ) + return model.to(device=device, dtype=dtype) + + +def _log_checkpoint_metadata(ckpt_path: Path) -> None: + """Log ``step``/``tokens_seen`` from the plain-JSON ``metadata.json`` if + present. Never reads ``train_state.pt`` (a pickle behind a UID-ownership + security gate); only the model weights from the ``.distcp`` shards are + needed for inference. + """ + meta_file = ckpt_path / "metadata.json" + if not meta_file.exists(): + return + try: + meta = json.loads(meta_file.read_text()) + except (json.JSONDecodeError, OSError) as exc: + logger.warning(f"Could not read {meta_file}: {exc}") + return + logger.info( + f"VLM checkpoint metadata: step={meta.get('step')}, tokens_seen={meta.get('tokens_seen')}" + ) + + +def _load_config(config_path: str) -> JobConfig: + config = load_config(config_path, cli_args=[]) + if not config.is_vlm: + raise ValueError( + f"{config_path!r} is not a VLM config (config.vlm is None); this evaluation " + f"path is VLM-only. Use scripts/eval.py for text-model loss/perplexity." + ) + return config + + +def _check_generative(vlm_config: VLMConfig) -> None: + """Fail fast (before building) on arches that cannot autoregressively generate.""" + if not vlm_config.is_generative: + raise ValueError( + f"VLM arch {vlm_config.arch!r} cannot be evaluated: its routing is non-causal " + f"and cannot autoregressively generate, but chat tasks are generation-only. " + f"Generation support for {vlm_config.arch!r} is a tracked model-side follow-up " + f"— contact the project owner." + ) + + +def _load_weights( + config: JobConfig, checkpoint: str, device: torch.device, dtype: torch.dtype +) -> VLMWrapper: + """Build a ``VLMWrapper`` and load DCP weights for single-process eval. + + Accepts either a run directory (resolved to its ``latest``/highest + ``step_N`` via ``resolve_resume_path``) or a specific checkpoint directory + (used as-is when ``resolve_resume_path`` finds nothing). DCP reshards on + load, so checkpoints saved under FSDP/PP load into the full model. + """ + ckpt_path = resolve_resume_path(checkpoint) or Path(checkpoint) + if not ckpt_path.exists(): + raise FileNotFoundError(f"Checkpoint path does not exist: {ckpt_path}") + if not ckpt_path.is_dir(): + raise NotADirectoryError(f"Checkpoint path is not a directory: {ckpt_path}") + model = _build_model(config, device, dtype) + model.eval() + + # Single-process DCP load: build the full (unsharded) model, then load the + # model shards into its state-dict. + state_dict = {"model": model.state_dict()} + dcp.load(state_dict, checkpoint_id=str(ckpt_path)) + model.load_state_dict(state_dict["model"]) + + _log_checkpoint_metadata(ckpt_path) + logger.info(f"Loaded VLM checkpoint from {ckpt_path}") + return model + + +# --------------------------------------------------------------------------- # +# Request rendering + preprocessing +# --------------------------------------------------------------------------- # + + +def _render_request( + messages: ChatMessages, video_config: VideoConfig | None +) -> tuple[list[Any], str]: + """Flatten one chat request into ``(frames, prompt_text)``. + + ``frames`` is an ordered list of visual frames; a single image is the + length-1 case and a decoded video is the multi-frame case. ``video_config`` + is the checkpoint's ``[video]`` config (``None`` for an image checkpoint) + and selects the mode: + + - **Image checkpoint** (``video_config is None``): exactly one image per + request; video content raises (an image model cannot evaluate video). + - **Video checkpoint**: exactly one video — decoded to frames via + ``video_io.decode_video_frames`` using the checkpoint's frame-sampling + policy — or, when no video is present, a single image treated as a + 1-frame clip (zero-padded to ``frames_per_clip`` downstream). + + Out-of-scope content (audio, multi-turn/few-shot, multi-image, multiple + videos, mixed image+video) raises ``NotImplementedError`` so the offending + task is surfaced rather than silently mishandled. Text content blocks are + concatenated in message order (newline-joined); role/turn structure is + intentionally discarded (see the module docstring on flattening). + """ + images, videos, audios = messages.extract_media() + if audios: + raise NotImplementedError( + "Audio evaluation is not implemented (image/video only). An audio request " + "reached the KempnerForge VLM adapter; report the task to the project owner." + ) + + roles = [message.role for message in messages.messages] + if any(role == "assistant" for role in roles) or roles.count("user") > 1: + raise NotImplementedError( + "Multi-turn / few-shot requests are not supported (single-turn, zero-shot " + "only). Report the task to the project owner." + ) + + prompt = "\n".join( + content.text + for message in messages.messages + for content in message.content + if content.type == "text" + ) + + if video_config is None: + # Image checkpoint: image-only, exactly one image per request. + if videos: + raise NotImplementedError( + "This is an image checkpoint (no [video] config) and cannot evaluate video. " + "Use a video checkpoint, or report the task to the project owner." + ) + if len(images) != 1: + raise NotImplementedError( + f"This adapter supports exactly one image per request, got {len(images)}. " + "Multi-image and text-only requests are out of scope; report the task to " + "the project owner." + ) + return images, prompt + + # Video checkpoint: exactly one visual — a video (decoded to frames) or a + # single image (treated as a 1-frame clip, zero-padded downstream). + if len(videos) > 1: + raise NotImplementedError( + f"Multiple videos per request are not supported, got {len(videos)}. " + "Report the task to the project owner." + ) + if videos: + if images: + raise NotImplementedError( + "Mixed image + video content in one request is not supported. " + "Report the task to the project owner." + ) + path = videos[0] + if not isinstance(path, str): + raise NotImplementedError( + f"Video content must be a local path string for decoding, got " + f"{type(path).__name__}. The task may pass clip boundaries or a URL; " + "report the task to the project owner." + ) + frames = decode_video_frames( + path, + fps=video_config.fps, + min_frames=video_config.min_frames, + max_frames=video_config.max_frames, + sampling_policy=video_config.sampling_policy, + ) + if not frames: + logger.warning( + f"No frames decoded from {path}; evaluating a zero clip (result unreliable)." + ) + return frames, prompt + if len(images) == 1: + # A single image on a video checkpoint: a 1-frame clip (zero-padded to + # frames_per_clip downstream), consistent with how training pads short clips. + return images, prompt + raise NotImplementedError( + f"A video-checkpoint request must carry exactly one video or one image, got " + f"{len(images)} images and no video. Multi-image and text-only requests are out " + "of scope; report the task to the project owner." + ) + + +def _to_pil(frame: Any) -> Any: + """Normalize one visual frame to a PIL ``Image``. + + lmms-eval delivers image content as either a ``PIL.Image`` or a path/URL + string. The training preprocessing (``pil_to_tensor``) is strict PIL-only, so + a string is opened here (copied so the file handle can close). This lets both + the image packer (``_frames_to_pixel_values``) and the video packer + (``frames_to_clip_tensor`` -> ``pil_to_tensor``) accept either form. + """ + if isinstance(frame, str): + from PIL import Image + + with Image.open(frame) as im: + return im.copy() + return frame + + +def _frames_to_pixel_values( + frames: list[Any], image_size: int, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """Convert an ordered list of frames to a ``(num_frames, 3, H, W)`` tensor. + + Reuses the training-time ``pil_to_tensor`` (resize + SigLIP-style + normalize) as the single source of truth. v1 passes a single image (a + length-1 list); the list shape is the seam for future video (ordered + frames). Each frame may be a ``PIL.Image`` or a path string (see ``_to_pil``). + """ + tensors = [ + pil_to_tensor(_to_pil(frame), image_size, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + for frame in frames + ] + return torch.stack(tensors, dim=0).to(device=device, dtype=dtype) + + +# --------------------------------------------------------------------------- # +# Generation +# --------------------------------------------------------------------------- # + + +class _ContextBudgetError(ValueError): + """No room for the prompt: max_new_tokens + image tokens exceed max_seq_len.""" + + +def _resolve_gen_kwargs(gen_kwargs: dict[str, Any], default_max_new_tokens: int) -> dict[str, Any]: + """Merge task ``gen_kwargs`` over the adapter's fallback defaults. + + Uses explicit ``is None`` checks (not ``x or default``) so an explicit falsy + task value — e.g. ``max_new_tokens=0`` or ``temperature=0.0`` (greedy) — is + honored rather than silently replaced by the default. ``top_p`` is the one + exception: a non-positive value is an ill-defined nucleus (an empty token + set), so ``top_p <= 0`` is treated as disabled (1.0) instead of passed + through to ``sample`` (where it would mask every token and raise). + """ + until = gen_kwargs.get("until") + if until is None: + until = [] + elif isinstance(until, str): + until = [until] + + mnt = gen_kwargs.get("max_new_tokens") + max_new_tokens = default_max_new_tokens if mnt is None else int(mnt) + temp = gen_kwargs.get("temperature") + temperature = 0.0 if temp is None else float(temp) + # An explicit do_sample=False forces greedy even if a temperature is given. + if not gen_kwargs.get("do_sample", temperature > 0): + temperature = 0.0 + + sampling = temperature > 0 + top_k = gen_kwargs.get("top_k") + top_p = gen_kwargs.get("top_p") + top_p = 1.0 if top_p is None else float(top_p) + + # Treat undefined `top_p` as diasbled + if top_p <= 0: + top_p = 1.0 + return { + "until": [u for u in until if u], + "max_new_tokens": max_new_tokens, + "temperature": temperature, + "top_k": (0 if top_k is None else int(top_k)) if sampling else 0, + "top_p": top_p if sampling else 1.0, + } + + +def _first_stop(text: str, until: list[str]) -> int | None: + """Index of the earliest occurrence of any stop string in ``text``.""" + cut: int | None = None + for stop in until: + idx = text.find(stop) + if idx != -1 and (cut is None or idx < cut): + cut = idx + return cut + + +@torch.inference_mode() +def _generate_batch( + model: VLMWrapper, + tokenizer: Any, + pixel_values: torch.Tensor, + prompt_ids: list[torch.Tensor], + resolved: dict[str, Any], + max_seq_len: int, + frame_mask: torch.Tensor | None = None, +) -> list[str]: + """Batched decode (no transformer KV cache); returns one continuation per request. + + Decodes ``B`` requests together (``pixel_values`` is ``(B, 3, H, W)`` or a + ``(B, F, 3, H, W)`` video clip, ``prompt_ids`` a list of ``B`` 1-D token + tensors; ``frame_mask`` is ``(B, F)`` bool for video, masking padded-frame + visual tokens from attention as in training). There is no transformer KV + cache and no vision cache: ``model(...)`` re-runs over the growing + **right-padded** batch each step, re-encoding the vision tower each time. + Right-padding matches the training + layout: the image prefix stays at positions ``0..n-1`` and text is contiguous + from ``n`` for every row (so image/text RoPE distances are consistent across + rows), and the trailing pads are causally masked, so a batched forward gives + each row the same real-position logits as decoding it alone. Each row's next + token is read at its own last real position; EOS / ``max_new_tokens`` / first + ``until`` match are tracked per row. ``B == 1`` reproduces the single-request + path exactly. + """ + until: list[str] = resolved["until"] + max_new_tokens: int = resolved["max_new_tokens"] + temperature: float = resolved["temperature"] + top_k: int = resolved["top_k"] + top_p: float = resolved["top_p"] + eos_id = tokenizer.eos_token_id + pad_id = resolve_pad_id(tokenizer) + device = pixel_values.device + batch_size = len(prompt_ids) + + # Length bound: image tokens (in-residual for JD/MoT; 0 for CA) + prompt + + # generated must fit the context. Reserve room for generation and left- + # truncate any over-budget prompt (per row). + num_image_tokens = model.num_image_tokens + prompt_budget = max_seq_len - num_image_tokens - max_new_tokens + if prompt_budget <= 0: + raise _ContextBudgetError( + f"max_new_tokens ({max_new_tokens}) + image tokens ({num_image_tokens}) leave no " + f"room for the prompt within max_seq_len ({max_seq_len}); lower --max-new-tokens." + ) + prompts: list[torch.Tensor] = [] + for ids in prompt_ids: + if ids.shape[0] > prompt_budget: + logger.warning( + f"Prompt ({ids.shape[0]}) + image tokens ({num_image_tokens}) + max_new_tokens " + f"({max_new_tokens}) exceeds max_seq_len ({max_seq_len}); left-truncating prompt " + f"to {prompt_budget} tokens. Severe truncation may distort results." + ) + ids = ids[-prompt_budget:] + prompts.append(ids) + + generated: list[list[int]] = [[] for _ in range(batch_size)] + done = [False] * batch_size + row_index = torch.arange(batch_size, device=device) + + for _ in range(max_new_tokens): + # Rebuild the right-padded batch from prompt + tokens generated so far. + seqs = [ + torch.cat([prompts[i], torch.tensor(generated[i], dtype=torch.long, device=device)]) + for i in range(batch_size) + ] + real_len = torch.tensor([s.shape[0] for s in seqs], device=device) + cur_max = int(real_len.max().item()) + input_ids = torch.full((batch_size, cur_max), pad_id, dtype=torch.long, device=device) + for i, s in enumerate(seqs): + input_ids[i, : s.shape[0]] = s + + logits, _ = model(pixel_values, input_ids, frame_mask=frame_mask) + # Each row's next-token logits sit at its own last real position (the + # output is already trimmed to text positions for JD/MoT; CA has no + # image prefix), not at [-1] (a pad for shorter rows). + next_logits = logits[row_index, real_len - 1] + next_tokens = sample(next_logits, temperature, top_k, top_p) + + for i in range(batch_size): + if done[i]: + continue + token_id = int(next_tokens[i].item()) + if eos_id is not None and token_id == eos_id: + done[i] = True + continue + generated[i].append(token_id) + if len(generated[i]) >= max_new_tokens: + done[i] = True + elif until: + text = tokenizer.decode(generated[i], skip_special_tokens=True) + if _first_stop(text, until) is not None: + done[i] = True + if all(done): + break + + outputs: list[str] = [] + for tokens in generated: + text = tokenizer.decode(tokens, skip_special_tokens=True) + cut = _first_stop(text, until) + outputs.append(text[:cut] if cut is not None else text) + return outputs + + +# --------------------------------------------------------------------------- # +# Adapter +# --------------------------------------------------------------------------- # + + +class KempnerForgeVLM(lmms): + """lmms-eval chat model over a KempnerForge ``VLMWrapper`` (see module docstring). + + Model args (parsed by the base ``create_from_arg_string`` from a + ``key=value,...`` string): + + - ``config`` (required): path to the KempnerForge TOML the checkpoint was + trained with. + - ``checkpoint`` (required): DCP checkpoint directory (a run dir or a + specific ``step_N`` dir). + - ``device`` (default ``"cuda"``), ``dtype`` (default: the checkpoint config's + ``train.param_dtype``; pass e.g. ``"float32"`` to override). + - ``batch_size`` (default ``1``): number of requests decoded together + (right-padded), grouped by gen_kwargs. + - ``max_new_tokens`` (default ``128``): fallback only; task ``gen_kwargs`` + override it. + """ + + is_simple = False + + def __init__( + self, + config: str, + checkpoint: str, + device: str = "cuda", + dtype: str | None = None, + batch_size: int | str = 1, + max_new_tokens: int | str = DEFAULT_MAX_NEW_TOKENS, + **kwargs: Any, + ) -> None: + super().__init__() + if kwargs: + logger.warning(f"Ignoring unsupported model_args: {sorted(kwargs)}") + + self._device = torch.device(device) + self._batch_size = int(batch_size) + self._default_max_new_tokens = int(max_new_tokens) + if self._batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {self._batch_size}") + if self._default_max_new_tokens < 1: + raise ValueError(f"max_new_tokens must be >= 1, got {self._default_max_new_tokens}") + self._config = _load_config(config) + assert self._config.vlm is not None # guaranteed by is_vlm; narrows for the type checker + # Default the compute dtype to what the checkpoint was trained at + # (config.train.param_dtype) unless an explicit dtype was passed. + self._dtype = self._config.train.param_dtype if dtype is None else _resolve_dtype(dtype) + self._arch = self._config.vlm.arch + # Fail fast on non-generative arches before building/loading the model. + _check_generative(self._config.vlm) + + # Video vs image mode is a property of the checkpoint's config. A video + # checkpoint ([video] config) fixes frames_per_clip and resizes frames to + # config.video.frame_size; an image checkpoint is a 1-frame clip at + # config.data.hf_image_size. + self._is_video = self._config.is_video + if self._config.video is not None: + self._frames_per_clip = self._config.video.max_frames + self._frame_size = self._config.video.frame_size + else: + self._frames_per_clip = 1 + self._frame_size = self._config.data.hf_image_size + + self._model = _load_weights(self._config, checkpoint, self._device, self._dtype) + self._tokenizer = build_tokenizer(self._config.data.tokenizer_path) + self._max_seq_len = self._config.model.max_seq_len + logger.info( + f"KempnerForgeVLM ready: arch={self._arch}, video={self._is_video}, " + f"frames_per_clip={self._frames_per_clip}, device={self._device}, " + f"dtype={self._dtype}, max_seq_len={self._max_seq_len}" + ) + + def generate_until(self, requests: list[Instance]) -> list[str]: + # Group requests by gen_kwargs (a batch must share decode params) and, + # within a group, sort by context length so similar-length prompts batch + # together (less padding). Collator.get_original restores request order. + def _collate(args: tuple[Any, ...]) -> int: + return -len(args[0]) if isinstance(args[0], str) else 0 + + re_ords = Collator( + [request.args for request in requests], + _collate, + group_fn=lambda args: args[2], # args[2] == gen_kwargs + grouping=True, + ) + results: list[str] = [] + pbar = tqdm(total=len(requests), disable=(self.rank != 0), desc="KempnerForge VLM") + for chunk in re_ords.get_batched(n=self._batch_size, batch_fn=None): + # Every request in the chunk shares gen_kwargs (index 2); resolve once. + resolved = _resolve_gen_kwargs(chunk[0][2], self._default_max_new_tokens) + # Per-request slots aligned to ``chunk``: ``None`` = filled by generation + # below; ``""`` = a request that failed to render/preprocess, isolated + # with a warning so one bad doc does not abort the whole run. + chunk_outputs: list[str | None] = [None] * len(chunk) + frames_batch: list[torch.Tensor] = [] + masks_batch: list[torch.Tensor] = [] + prompt_ids: list[torch.Tensor] = [] + for slot, args in enumerate(chunk): + try: + # Chat 6-tuple: (context, doc_to_messages, gen_kwargs, doc_id, task, split). + doc = self.task_dict[args[4]][args[5]][args[3]] + messages = ChatMessages(messages=args[1](doc)) + frames, prompt = _render_request(messages, self._config.video) + # lmms-eval may deliver image content as a path/URL string; + # normalize to PIL so both the image and video packers (strict + # pil_to_tensor) accept it. + frames = [_to_pil(f) for f in frames] + # Mirror training tokenization: no chat template, no + # placeholder, add_special_tokens=False (images go via pixel_values). + token_ids = self._tokenizer(prompt, add_special_tokens=False)["input_ids"] + if not token_ids: + # No text to condition on. Image-prefix positions are trained + # with -100 labels (and trimmed by output_slice), so there is + # no valid position to predict an image-only first token. + raise ValueError("empty prompt after flattening (no text content)") + if self._is_video: + # Fixed (frames_per_clip, 3, H, W) clip + per-frame validity + # mask, zero-padded — identical to training. The mask hides + # padded-frame visual tokens from attention. + clip, fmask = frames_to_clip_tensor( + frames, max_frames=self._frames_per_clip, frame_size=self._frame_size + ) + pixels = clip.to(device=self._device, dtype=self._dtype) + mask = fmask.to(device=self._device) + else: + pixels = _frames_to_pixel_values( + frames, self._frame_size, self._device, self._dtype + ) + mask = None + prompt_tensor = torch.tensor(token_ids, dtype=torch.long, device=self._device) + except Exception as exc: + # Isolate a bad request (unsupported content, decode failure, empty + # prompt, ...) so the remaining requests still score. + logger.warning( + f"Skipping request (task={args[4]}, doc_id={args[3]}): " + f"{type(exc).__name__}: {exc}" + ) + chunk_outputs[slot] = "" + continue + # Commit atomically: only reached when every step above succeeded, so a + # mid-request failure never leaves a partial entry in these lists. + frames_batch.append(pixels) + if mask is not None: + masks_batch.append(mask) + prompt_ids.append(prompt_tensor) + + if prompt_ids: + # Video: (B, F, 3, H, W) via stack (each request is one F-frame clip), + # with a (B, F) frame mask. Image: (B, 3, H, W) via cat. cat on video + # would fold frames into the batch and trip the frames-per-clip check. + if self._is_video: + pixel_values = torch.stack(frames_batch, dim=0) + frame_mask = torch.stack(masks_batch, dim=0) + else: + pixel_values = torch.cat(frames_batch, dim=0) + frame_mask = None + try: + gen_outputs = _generate_batch( + self._model, + self._tokenizer, + pixel_values, + prompt_ids, + resolved, + self._max_seq_len, + frame_mask=frame_mask, + ) + except _ContextBudgetError as exc: + # One task's gen_kwargs over-budgets the context; skip its requests + # (they all share gen_kwargs) rather than aborting the whole run. + logger.warning( + f"Skipping {len(prompt_ids)} request(s) for task {chunk[0][4]}: {exc}" + ) + gen_outputs = [""] * len(prompt_ids) + else: + gen_outputs = [] + # Scatter generated continuations back into the surviving (None) slots, + # preserving alignment with ``chunk`` (skipped slots keep their ""). + gen_iter = iter(gen_outputs) + outputs = [o if o is not None else next(gen_iter) for o in chunk_outputs] + for args, output in zip(chunk, outputs, strict=True): + results.append(output) + self.cache_hook.add_partial("generate_until", (args[0], args[2]), output) + pbar.update(len(chunk)) + pbar.close() + return re_ords.get_original(results) + + def loglikelihood(self, requests: list[Instance]) -> list[tuple[float, bool]]: + raise NotImplementedError( + "KempnerForgeVLM is a generation-only chat model; loglikelihood is not supported. " + "Standard multiple-choice VLM benchmarks run as generate_until tasks in lmms-eval." + ) + + def generate_until_multi_round(self, requests: list[Instance]) -> list[str]: + raise NotImplementedError("KempnerForgeVLM does not support multi-round generation in v1.") diff --git a/kempnerforge/eval/vlm/manifest.py b/kempnerforge/eval/vlm/manifest.py new file mode 100644 index 0000000..3341509 --- /dev/null +++ b/kempnerforge/eval/vlm/manifest.py @@ -0,0 +1,19 @@ +# pyright: reportMissingImports=false +# ^ lmms-eval is an optional, undeclared dependency; see adapter.py for why this +# file-level directive is used instead of an inline ignore. +"""lmms-eval registration manifest for the KempnerForge VLM adapter. + +``MANIFEST`` is discovered by lmms-eval through the ``lmms_eval.models`` entry +point declared in ``pyproject.toml`` (``kempnerforge_vlm = +"kempnerforge.eval.vlm.manifest:MANIFEST"``). Only ``chat_class_path`` is set: +the adapter is chat-only (``is_simple = False``), so resolution stays correct +even under ``force_simple``. The entry point is metadata only and does not make +lmms-eval a runtime dependency of KempnerForge. +""" + +from lmms_eval.models.registry_v2 import ModelManifest + +MANIFEST = ModelManifest( + model_id="kempnerforge_vlm", + chat_class_path="kempnerforge.eval.vlm.adapter.KempnerForgeVLM", +) diff --git a/kempnerforge/model/vlm.py b/kempnerforge/model/vlm.py index c6ca8f9..d6dfd74 100644 --- a/kempnerforge/model/vlm.py +++ b/kempnerforge/model/vlm.py @@ -180,7 +180,7 @@ def prepare( self, wrapper: VLMWrapper, pixel_values: torch.Tensor, - input_ids: torch.Tensor, + input_ids: torch.Tensor, # noqa: ARG002 frame_mask: torch.Tensor | None = None, ) -> ModalityContext: img_embeds = _project_visual_features(wrapper, pixel_values) diff --git a/pyproject.toml b/pyproject.toml index acfd191..b3cb367 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,9 @@ dependencies = [ "torchao>=0.17.0", ] +[project.entry-points."lmms_eval.models"] +kempnerforge_vlm = "kempnerforge.eval.vlm.manifest:MANIFEST" + [dependency-groups] dev = [ "pyright[nodejs]>=1.1.408", diff --git a/scripts/vlm_eval_harness.py b/scripts/vlm_eval_harness.py new file mode 100644 index 0000000..48aeca0 --- /dev/null +++ b/scripts/vlm_eval_harness.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Run lmms-eval benchmarks on a KempnerForge VLM checkpoint. + +Evaluates a VLM checkpoint directly via the +``kempnerforge_vlm`` lmms-eval chat-model adapter, on the standard image +benchmarks lmms-eval implements as ``generate_until`` tasks (MMMU, MMBench, +ScienceQA, SEED, AI2D, ...). + +Requirements (lmms-eval is an OPTIONAL, separately-installed dependency, exactly +like lm-eval for text evaluation): + + uv pip install lmms-eval + +v1 is single-GPU and image-only; MoMa checkpoints are not supported (see +docs/how-to/run-vlm-evaluation.md). On clusters where importing lmms-eval's +evaluator fails with ``GLIBCXX_... not found``, put a newer libstdc++ on the +library path (e.g. ``LD_LIBRARY_PATH=/lib``). + +Usage: + uv run python scripts/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val \ + --output results/vlm_step_10000.json + + # Quick partial run (4 examples per task) + uv run python scripts/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val,mmbench_en_dev \ + --limit 4 +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from pathlib import Path + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + + +def _limit_type(value: str) -> int | float: + """Per-task example cap: an integer count, or a fraction < 1.0.""" + parsed = float(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("--limit must be > 0") + if parsed < 1.0: + return parsed + if parsed.is_integer(): + return int(parsed) + raise argparse.ArgumentTypeError("--limit must be an integer count, or a fraction < 1.0") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Run lmms-eval on a KempnerForge VLM checkpoint", + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="KempnerForge TOML the checkpoint was trained with", + ) + parser.add_argument( + "--checkpoint", type=str, required=True, help="DCP checkpoint dir (run dir or step_N dir)" + ) + # No default task suite: the representative default benchmark set is an open + # decision; --tasks is required until one is provided. + parser.add_argument( + "--tasks", type=str, required=True, help="Comma-separated lmms-eval task names" + ) + parser.add_argument("--output", type=str, default=None, help="Output JSON file path") + parser.add_argument( + "--limit", + type=_limit_type, + default=None, + help="Cap examples per task (int count, or <1.0 fraction); for quick partial runs", + ) + parser.add_argument("--device", type=str, default="cuda", help="Device (default: cuda)") + parser.add_argument( + "--dtype", + type=str, + default=None, + help="Model dtype; default: the checkpoint config's train.param_dtype", + ) + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Requests decoded together (grouped by gen_kwargs)", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=128, + help="Fallback max new tokens; task gen_kwargs override it (default: 128)", + ) + args = parser.parse_args() + + # lmms-eval is optional and undeclared; import lazily with a helpful error. + try: + from lmms_eval.evaluator import simple_evaluate + except ImportError as exc: + logger.error( + "Could not import lmms-eval's simple_evaluate (%s).\n" + "lmms-eval is an optional dependency; install it with: uv pip install lmms-eval\n" + "If this is a 'GLIBCXX_... not found' error, put a newer libstdc++ on the library " + "path (e.g. LD_LIBRARY_PATH=/lib); see docs/how-to/run-vlm-evaluation.md.", + exc, + ) + sys.exit(1) + + logger.info(f"Running lmms-eval: tasks={args.tasks}, checkpoint={args.checkpoint}") + + model_args = ( + f"config={args.config},checkpoint={args.checkpoint},max_new_tokens={args.max_new_tokens}" + ) + # Only pass dtype when explicitly set; otherwise the adapter defaults it from + # the checkpoint config (train.param_dtype). + if args.dtype is not None: + model_args += f",dtype={args.dtype}" + results = simple_evaluate( + model="kempnerforge_vlm", + model_args=model_args, + tasks=args.tasks.split(","), + device=args.device, + batch_size=args.batch_size, + limit=args.limit, + ) + + # --- Print results --- + print(f"\n{'=' * 60}") + print("lmms-eval Results") + print(f"{'=' * 60}") + if results is not None and "results" in results: + for task_name, task_results in sorted(results["results"].items()): + print(f"\n {task_name}:") + for metric, value in sorted(task_results.items()): + if isinstance(value, float): + print(f" {metric}: {value:.4f}") + elif metric != "alias": + print(f" {metric}: {value}") + print(f"{'=' * 60}\n") + + # --- Save results --- + if args.output: + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(results, f, indent=2, default=str) + logger.info(f"Results saved to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/conftest.py b/tests/conftest.py index b651c4a..629a736 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,13 +7,17 @@ import torch from kempnerforge.config.schema import ( + AdapterConfig, CheckpointConfig, JobConfig, ModelConfig, OptimizerConfig, SchedulerConfig, TrainConfig, + VisionEncoderConfig, + VLMConfig, ) +from kempnerforge.config.video import VideoConfig # --------------------------------------------------------------------------- # CLI flags for opt-in test suites @@ -86,6 +90,54 @@ def tiny_job_config(tmp_path) -> JobConfig: ) +# --------------------------------------------------------------------------- +# VLM configs / model +# --------------------------------------------------------------------------- + + +@pytest.fixture +def tiny_vlm_configs() -> tuple[ModelConfig, VisionEncoderConfig, AdapterConfig, VLMConfig]: + """Tiny Joint-Decoder VLM configs for CPU tests (random vision encoder).""" + return ( + ModelConfig(dim=64, n_layers=2, n_heads=4, vocab_size=256, max_seq_len=64), + VisionEncoderConfig(type="random", feature_dim=96, num_tokens=8), + AdapterConfig(), + VLMConfig(max_text_len=32), + ) + + +@pytest.fixture +def tiny_vlm_wrapper(tiny_vlm_configs): + """A tiny CPU ``VLMWrapper`` built from ``tiny_vlm_configs`` (no checkpoint).""" + from kempnerforge.model.vlm import build_vlm_wrapper + + return build_vlm_wrapper(*tiny_vlm_configs).eval() + + +@pytest.fixture +def tiny_video_configs( + tiny_vlm_configs, +) -> tuple[ModelConfig, VisionEncoderConfig, AdapterConfig, VLMConfig, VideoConfig]: + """Tiny configs for a video VLM checkpoint (``frames_per_clip == 2``). + + Reuses the image VLM configs and adds a ``[video]`` section sized so the + visual budget (2 frames x 8 tokens = 16) plus ``max_text_len`` (32) fits the + tiny ``max_seq_len`` (64). + """ + mc, vc, ac, lc = tiny_vlm_configs + video = VideoConfig(max_frames=2, min_frames=1, frame_size=16) + return (mc, vc, ac, lc, video) + + +@pytest.fixture +def tiny_video_vlm_wrapper(tiny_video_configs): + """A tiny CPU video ``VLMWrapper`` (``frames_per_clip == 2``; no checkpoint).""" + from kempnerforge.model.vlm import build_vlm_wrapper + + mc, vc, ac, lc, video = tiny_video_configs + return build_vlm_wrapper(mc, vc, ac, lc, frames_per_clip=video.max_frames).eval() + + # --------------------------------------------------------------------------- # Device # --------------------------------------------------------------------------- diff --git a/tests/integration/test_lmms_eval_contract.py b/tests/integration/test_lmms_eval_contract.py new file mode 100644 index 0000000..f1d1746 --- /dev/null +++ b/tests/integration/test_lmms_eval_contract.py @@ -0,0 +1,113 @@ +"""Contract tests pinning the real ``lmms_eval`` API to what the unit-test fakes assume. + +The VLM-eval unit tests run against an in-repo fake ``lmms_eval`` +(``tests/unit/eval/vlm/_fake_lmms_eval.py``) so they execute in CI without the optional, +undeclared ``lmms-eval`` dependency. These tests are the fidelity net: they exercise the +*real* package and fail loudly if its API drifts from the fakes (in which case update the +fakes — and likely the adapter). They run wherever real lmms-eval is installed (locally, the +manual ``gpu-tests`` CI job) and skip otherwise. + +Also verifies the ``pyproject.toml`` ``lmms_eval.models`` entry point resolves to the adapter. +""" + +from __future__ import annotations + +import pytest + +lmms_eval = pytest.importorskip("lmms_eval") +if getattr(lmms_eval, "__file__", None) is None: + # The unit fake (a bare module object) is active; contract tests need the real package. + pytest.skip( + "fake lmms_eval is active; skipping real-package contract tests", allow_module_level=True + ) + +from lmms_eval.api.instance import Instance # noqa: E402 +from lmms_eval.api.model import lmms # noqa: E402 +from lmms_eval.models import get_model # noqa: E402 +from lmms_eval.models.registry_v2 import ModelManifest # noqa: E402 +from lmms_eval.protocol import ChatMessages # noqa: E402 +from lmms_eval.utils import Collator # noqa: E402 + +from kempnerforge.eval.vlm.adapter import KempnerForgeVLM # noqa: E402 + + +def test_entrypoint_resolves_to_adapter(): + """The pyproject ``lmms_eval.models`` entry point resolves ``kempnerforge_vlm``.""" + assert get_model("kempnerforge_vlm") is KempnerForgeVLM + + +class TestChatMessagesContract: + def test_extract_media_and_object_model(self): + messages = ChatMessages( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image", "url": "IMG"}, + ], + } + ] + ) + images, videos, audios = messages.extract_media() + assert images == ["IMG"] + assert videos == [] and audios == [] + message = messages.messages[0] + assert message.role == "user" + assert message.content[0].type == "text" and message.content[0].text == "hello" + assert message.content[1].type == "image" and message.content[1].url == "IMG" + + +class TestCollatorContract: + def test_group_sort_batch_and_restore(self): + # Mirrors the adapter: group_fn -> gen_kwargs dict, sort_fn -> -len(context). + arr = [ + ("longest-context", None, {"max_new_tokens": 4}), + ("c", None, {"max_new_tokens": 4}), + ("mid-ctx", None, {"max_new_tokens": 8}), + ] + col = Collator(arr, lambda a: -len(a[0]), group_fn=lambda a: a[2], grouping=True) + flat = [item for batch in col.get_batched(n=2, batch_fn=None) for item in batch] + assert sorted(map(id, flat)) == sorted(map(id, arr)) # all items present + assert col.get_original(flat) == arr # original request order restored + + +class TestModelManifestContract: + def test_fields_and_validation(self): + manifest = ModelManifest(model_id="x", chat_class_path="a.b.C") + assert manifest.model_id == "x" + assert manifest.chat_class_path == "a.b.C" + assert manifest.simple_class_path is None + with pytest.raises(ValueError): + ModelManifest(model_id="y") # neither class path -> rejected + + +class TestLmmsBaseContract: + def test_single_process_defaults(self): + class _Probe(lmms): + is_simple = False + + def loglikelihood(self, requests): + raise NotImplementedError + + def generate_until(self, requests): + raise NotImplementedError + + def generate_until_multi_round(self, requests): + raise NotImplementedError + + probe = _Probe() + assert probe.rank == 0 and probe.world_size == 1 + assert probe.task_dict == {} + probe.cache_hook.add_partial("generate_until", ("ctx", {}), "out") # no-op, must not raise + + +class TestInstanceContract: + def test_args_returns_arguments_tuple(self): + inst = Instance( + request_type="generate_until", + arguments=("ctx", None, {}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + assert inst.args == ("ctx", None, {}, "d0", "t", "test") diff --git a/tests/integration/test_vlm_eval.py b/tests/integration/test_vlm_eval.py new file mode 100644 index 0000000..2d74ede --- /dev/null +++ b/tests/integration/test_vlm_eval.py @@ -0,0 +1,208 @@ +"""Integration tests for the KempnerForge VLM lmms-eval adapter. + +Three tests, all skipped when lmms-eval is absent (optional, undeclared dep): + +1. ``test_dcp_roundtrip_generate_until`` — self-contained and CPU-only: builds a + tiny VLM, saves it via DCP, then loads it back through ``KempnerForgeVLM`` + and runs ``generate_until`` on a synthetic single-image request. Exercises + the real DCP load path and the full request -> render -> preprocess -> + decode -> collect pipeline without a GPU, a real checkpoint, or the network. + +2. ``test_dcp_roundtrip_video_generate_until`` — the video analogue: a tiny + ``frames_per_clip == 2`` video VLM, decoder stubbed, validating the video + build and the ``(B, F, 3, H, W)`` clip path on CPU. + +3. ``test_real_task_via_simple_evaluate`` — opt-in (set ``KF_VLM_EVAL_CONFIG`` + and ``KF_VLM_EVAL_CHECKPOINT``): runs a small ``--limit`` slice of a real + lmms-eval ``generate_until`` task through ``simple_evaluate`` against a real + checkpoint. Point it at a video config + video task to cover real video. + Intended for a GPU node; skipped by default. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed.checkpoint as dcp +from PIL import Image + +lmms_eval = pytest.importorskip("lmms_eval") +if getattr(lmms_eval, "__file__", None) is None: + # The unit conftest injects a fake lmms_eval (a bare module with no __file__); + # these DCP-roundtrip tests need the real package, so skip rather than bind the + # fake (mirrors tests/integration/test_lmms_eval_contract.py). + pytest.skip("fake lmms_eval is active; skipping real-package tests", allow_module_level=True) + +from lmms_eval.api.instance import Instance # noqa: E402 + +from kempnerforge.config.data import DataConfig # noqa: E402 +from kempnerforge.config.schema import JobConfig # noqa: E402 +from kempnerforge.eval.vlm.adapter import KempnerForgeVLM # noqa: E402 +from kempnerforge.model.vlm import build_vlm_wrapper # noqa: E402 + + +class _MockTokenizer: + pad_token_id = 0 + eos_token_id = None + + def __call__(self, text: str, add_special_tokens: bool = False) -> dict[str, list[int]]: + del add_special_tokens + return {"input_ids": [(ord(c) % 254) + 1 for c in text]} + + def decode(self, ids: list[int], skip_special_tokens: bool = True) -> str: + del skip_special_tokens + return " ".join(str(int(i)) for i in ids) + + +def test_dcp_roundtrip_generate_until(tmp_path, tiny_vlm_configs, monkeypatch): + mc, vc, ac, lc = tiny_vlm_configs + + # Build a tiny VLM and save it as a DCP checkpoint (single process). + torch.manual_seed(0) + model = build_vlm_wrapper(mc, vc, ac, lc).eval() + ckpt_dir = tmp_path / "step_0" + dcp.save({"model": model.state_dict()}, checkpoint_id=str(ckpt_dir)) + + # A matching JobConfig the loader will rebuild from. Patch _load_config to + # return it (config TOML parsing is covered by the config tests) and patch + # the tokenizer builder to avoid an HF download. + job_config = JobConfig( + model=mc, vision_encoder=vc, adapter=ac, vlm=lc, data=DataConfig(tokenizer_path="mock") + ) + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._load_config", lambda _path: job_config) + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.build_tokenizer", lambda _path: _MockTokenizer() + ) + + vlm = KempnerForgeVLM( + config="ignored", checkpoint=str(ckpt_dir), device="cpu", dtype="float32", batch_size=2 + ) + + # Two synthetic single-image requests with different prompt lengths, decoded + # as one right-padded batch (batch_size=2), mirroring the chat 6-tuple. + img = Image.new("RGB", (8, 8), color=(120, 120, 120)) + + def doc_to_messages(doc): + return [ + { + "role": "user", + "content": [{"type": "text", "text": doc["q"]}, {"type": "image", "url": img}], + } + ] + + vlm.task_dict = { + "t": { + "test": { + "d0": {"q": "What color is this?"}, + "d1": {"q": "Describe this picture in a few words please."}, + } + } + } + instances = [ + Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, doc_id in enumerate(["d0", "d1"]) + ] + + outputs = vlm.generate_until(instances) + assert isinstance(outputs, list) and len(outputs) == 2 + assert all(isinstance(o, str) for o in outputs) + assert all(len(o.split()) == 3 for o in outputs) # greedy emits exactly max_new_tokens + + +def test_dcp_roundtrip_video_generate_until(tmp_path, tiny_video_configs, monkeypatch): + """Video analogue of the image roundtrip: a tiny video VLM (frames_per_clip == + 2) is DCP-saved and reloaded through ``KempnerForgeVLM``, then evaluated on + synthetic video requests with the decoder stubbed (no ``av`` / no real clip). + Exercises the real DCP load, the ``frames_per_clip`` build, and the + render -> decode -> (B, F, 3, H, W) stack -> forward -> decode video path.""" + mc, vc, ac, lc, video = tiny_video_configs + + torch.manual_seed(0) + model = build_vlm_wrapper(mc, vc, ac, lc, frames_per_clip=video.max_frames).eval() + ckpt_dir = tmp_path / "step_0" + dcp.save({"model": model.state_dict()}, checkpoint_id=str(ckpt_dir)) + + job_config = JobConfig( + model=mc, + vision_encoder=vc, + adapter=ac, + vlm=lc, + video=video, + data=DataConfig(tokenizer_path="mock"), + ) + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._load_config", lambda _path: job_config) + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.build_tokenizer", lambda _path: _MockTokenizer() + ) + # Stub video decode so the test needs no av / no real video file. + frames = [Image.new("RGB", (8, 8), color=(120, 120, 120)) for _ in range(2)] + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.decode_video_frames", lambda path, **kw: frames + ) + + vlm = KempnerForgeVLM( + config="ignored", checkpoint=str(ckpt_dir), device="cpu", dtype="float32", batch_size=2 + ) + assert vlm._is_video and vlm._frames_per_clip == 2 + + def doc_to_messages(doc): + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": doc["q"]}, + {"type": "video", "url": doc["v"]}, + ], + } + ] + + vlm.task_dict = { + "t": { + "test": { + "d0": {"q": "What happens in this clip?", "v": "a.mp4"}, + "d1": {"q": "Describe the video briefly please.", "v": "b.mp4"}, + } + } + } + instances = [ + Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, doc_id in enumerate(["d0", "d1"]) + ] + + outputs = vlm.generate_until(instances) + assert isinstance(outputs, list) and len(outputs) == 2 + assert all(isinstance(o, str) for o in outputs) + assert all(len(o.split()) == 3 for o in outputs) # greedy emits exactly max_new_tokens + + +@pytest.mark.skipif( + not os.environ.get("KF_VLM_EVAL_CONFIG") or not os.environ.get("KF_VLM_EVAL_CHECKPOINT"), + reason="set KF_VLM_EVAL_CONFIG and KF_VLM_EVAL_CHECKPOINT to run against a real checkpoint", +) +def test_real_task_via_simple_evaluate(): + from lmms_eval.evaluator import simple_evaluate + + config = os.environ["KF_VLM_EVAL_CONFIG"] + checkpoint = os.environ["KF_VLM_EVAL_CHECKPOINT"] + task = os.environ.get("KF_VLM_EVAL_TASK", "mmmu_val") + + results = simple_evaluate( + model="kempnerforge_vlm", + model_args=f"config={config},checkpoint={checkpoint}", + tasks=[task], + limit=2, + ) + assert results is not None + assert "results" in results and task in results["results"] diff --git a/tests/unit/eval/__init__.py b/tests/unit/eval/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/eval/vlm/__init__.py b/tests/unit/eval/vlm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/eval/vlm/_fake_lmms_eval.py b/tests/unit/eval/vlm/_fake_lmms_eval.py new file mode 100644 index 0000000..88fb8a3 --- /dev/null +++ b/tests/unit/eval/vlm/_fake_lmms_eval.py @@ -0,0 +1,237 @@ +"""Dependency-free fakes for the ``lmms_eval`` API surface the VLM adapter uses. + +``lmms-eval`` is an optional, *undeclared* dependency, so +``kempnerforge.eval.vlm.adapter`` / ``manifest`` import it at module top and cannot be +imported without it. ``conftest.py`` injects these fakes into ``sys.modules`` so the unit +tests run (and contribute coverage) in CI, where ``lmms-eval`` is absent. The fakes +reproduce ONLY the behavior the adapter relies on; their fidelity to the real package is +pinned by the gated contract test in ``tests/integration/`` (``test_lmms_eval_contract``). + +Verified against the installed ``lmms_eval`` source: + +- ``protocol.ChatMessages`` (pydantic): ``.messages[].role`` / ``.content[].type``; text + blocks carry ``.text``, image/video/audio blocks carry ``.url``. ``extract_media()`` + returns ``(images, videos, audios)`` of the ``.url`` payloads. +- ``utils.Collator``: groups by ``group_fn(item)`` (a dict; key built from sorted items + with list values tupled), sorts each group by ``sort_fn(item)``, ``get_batched(n)`` + yields lists of ``<= n`` while recording reorder indices, ``get_original`` inverts them. +- ``models.registry_v2.ModelManifest``: frozen dataclass requiring >= 1 class path. +- ``api.model.lmms``: base sets ``_rank=0/_world_size=1/cache_hook/task_dict`` and exposes + ``rank``/``world_size`` properties. +- ``api.instance.Instance``: dataclass exposing ``.args`` (the arguments tuple). +""" + +from __future__ import annotations + +import collections +import dataclasses +import types +from collections.abc import Callable +from typing import Any + +# --------------------------------------------------------------------------- # +# lmms_eval.protocol.ChatMessages +# --------------------------------------------------------------------------- # + + +class _Content: + def __init__(self, block: dict) -> None: + self.type = block["type"] + if self.type == "text": + self.text = block["text"] + else: + self.url = block["url"] + + +class _Message: + def __init__(self, message: dict) -> None: + self.role = message["role"] + self.content = [_Content(block) for block in message["content"]] + + +class ChatMessages: + def __init__(self, messages: list[dict]) -> None: + self.messages = [_Message(message) for message in messages] + + def extract_media(self) -> tuple[list[Any], list[Any], list[Any]]: + images: list[Any] = [] + videos: list[Any] = [] + audios: list[Any] = [] + for message in self.messages: + for content in message.content: + if content.type == "image": + images.append(content.url) + elif content.type == "video": + videos.append(content.url) + elif content.type == "audio": + audios.append(content.url) + return images, videos, audios + + +# --------------------------------------------------------------------------- # +# lmms_eval.utils.Collator +# --------------------------------------------------------------------------- # + + +def _hashable_key(group: dict) -> tuple: + # Mirror the real Collator: list/tuple values are tupled so the dict is hashable. + return tuple( + (key, tuple(value) if isinstance(value, (list, tuple)) else value) + for key, value in sorted(group.items()) + ) + + +class Collator: + def __init__( + self, + arr: list, + sort_fn: Callable, + group_fn: Callable = lambda x: x[1], + grouping: bool = False, + ) -> None: + self._sort_fn = sort_fn + self.size = len(arr) + self.reorder_indices: list[int] = [] + indexed = list(enumerate(arr)) + if grouping: + groups: dict[Any, list] = collections.OrderedDict() + for idx, item in indexed: + groups.setdefault(_hashable_key(group_fn(item)), []).append((idx, item)) + self._groups = list(groups.values()) + else: + self._groups = [indexed] + + def get_batched(self, n: int = 1, batch_fn: Callable | None = None): + del batch_fn # the adapter always passes None; size-by-n only + for group in self._groups: + ordered = sorted(group, key=lambda pair: self._sort_fn(pair[1])) + self.reorder_indices.extend(idx for idx, _ in ordered) + items = [item for _, item in ordered] + for start in range(0, len(items), n): + yield items[start : start + n] + + def get_original(self, newarr: list) -> list: + res: list[Any] = [None] * self.size + for ind, value in zip(self.reorder_indices, newarr, strict=True): + res[ind] = value + return res + + def __len__(self) -> int: + return self.size + + +# --------------------------------------------------------------------------- # +# lmms_eval.models.registry_v2.ModelManifest +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass(frozen=True) +class ModelManifest: + model_id: str + simple_class_path: str | None = None + chat_class_path: str | None = None + aliases: tuple = () + + def __post_init__(self) -> None: + if self.simple_class_path is None and self.chat_class_path is None: + raise ValueError(f"ModelManifest('{self.model_id}') requires at least one class path") + + +# --------------------------------------------------------------------------- # +# lmms_eval.api.model.lmms +# --------------------------------------------------------------------------- # + + +class _CacheHook: + def add_partial(self, attr: str, request: Any, result: Any) -> None: + pass + + +class lmms: # noqa: N801 — mirrors the real lowercase class name + is_simple: bool = True + + def __init__(self) -> None: + self._rank = 0 + self._world_size = 1 + self.cache_hook = _CacheHook() + self.task_dict: dict = {} + + @property + def rank(self) -> int: + return self._rank + + @property + def world_size(self) -> int: + return self._world_size + + +# --------------------------------------------------------------------------- # +# lmms_eval.api.instance.Instance +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass +class Instance: + request_type: str + arguments: tuple + idx: int + metadata: Any = dataclasses.field(default_factory=dict) + resps: list = dataclasses.field(default_factory=list) + filtered_resps: dict = dataclasses.field(default_factory=dict) + + def __post_init__(self) -> None: + meta = self.metadata or {} + self.task_name = meta.get("task") + self.doc_id = meta.get("doc_id") + self.repeats = meta.get("repeats") + + @property + def args(self) -> tuple: + return self.arguments if isinstance(self.arguments, tuple) else (self.arguments,) + + +# --------------------------------------------------------------------------- # +# Module tree assembly +# --------------------------------------------------------------------------- # + + +def build_modules() -> dict[str, types.ModuleType]: + """Build the fake ``lmms_eval`` module tree keyed by dotted import name. + + Submodules are also set as attributes on their parents so both ``import + a.b.c`` and ``a.b.c`` attribute access resolve once these are in ``sys.modules``. + """ + + def _mod(name: str, **attrs: Any) -> types.ModuleType: + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + return module + + root = _mod("lmms_eval") + api = _mod("lmms_eval.api") + api_model = _mod("lmms_eval.api.model", lmms=lmms, CacheHook=_CacheHook) + api_instance = _mod("lmms_eval.api.instance", Instance=Instance) + protocol = _mod("lmms_eval.protocol", ChatMessages=ChatMessages) + utils = _mod("lmms_eval.utils", Collator=Collator) + models = _mod("lmms_eval.models") + registry_v2 = _mod("lmms_eval.models.registry_v2", ModelManifest=ModelManifest) + + root.api = api + root.protocol = protocol + root.utils = utils + root.models = models + api.model = api_model + api.instance = api_instance + models.registry_v2 = registry_v2 + + return { + "lmms_eval": root, + "lmms_eval.api": api, + "lmms_eval.api.model": api_model, + "lmms_eval.api.instance": api_instance, + "lmms_eval.protocol": protocol, + "lmms_eval.utils": utils, + "lmms_eval.models": models, + "lmms_eval.models.registry_v2": registry_v2, + } diff --git a/tests/unit/eval/vlm/conftest.py b/tests/unit/eval/vlm/conftest.py new file mode 100644 index 0000000..3ad94e8 --- /dev/null +++ b/tests/unit/eval/vlm/conftest.py @@ -0,0 +1,44 @@ +"""Hermetic fake ``lmms_eval`` for the VLM-eval unit tests. + +``lmms-eval`` is an optional, undeclared dependency, so +``adapter.py``/``manifest.py`` cannot be imported without it and these tests would skip in +CI (no coverage). This conftest installs a faithful in-repo fake (``_fake_lmms_eval``) into +``sys.modules`` at import time so the tests always run and exercise our code. The fake is +installed unconditionally (hermetic): unit-test behavior is identical with or without real +lmms-eval present. Real-package fidelity is pinned separately by the gated contract test in +``tests/integration/`` (``test_lmms_eval_contract``). + +CI runs ``tests/unit/`` and ``tests/integration/`` as separate jobs, so the injected fake +never reaches the real-lmms-eval integration tests. For a combined local ``pytest tests/`` +run, ``tests/integration/`` is collected before ``tests/unit/`` (default ordering) and binds +the real package first; the saved originals are restored at session end. +""" + +from __future__ import annotations + +import sys + +import pytest + +from . import _fake_lmms_eval + +_ADAPTER_MODULES = ("kempnerforge.eval.vlm.adapter", "kempnerforge.eval.vlm.manifest") +_FAKE_MODULES = _fake_lmms_eval.build_modules() +_MANAGED = (*_FAKE_MODULES.keys(), *_ADAPTER_MODULES) +_SAVED = {name: sys.modules.get(name) for name in _MANAGED} + +# Install the fakes and evict any real-bound adapter/manifest so they re-import against the +# fakes when the test modules import them. +sys.modules.update(_FAKE_MODULES) +for _name in _ADAPTER_MODULES: + sys.modules.pop(_name, None) + + +def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: + del session, exitstatus + for name in _MANAGED: + saved = _SAVED[name] + if saved is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = saved diff --git a/tests/unit/eval/vlm/test_adapter.py b/tests/unit/eval/vlm/test_adapter.py new file mode 100644 index 0000000..2487607 --- /dev/null +++ b/tests/unit/eval/vlm/test_adapter.py @@ -0,0 +1,1133 @@ +"""CPU unit tests for the KempnerForge VLM lmms-eval adapter. + +A faithful fake ``lmms_eval`` is injected by ``conftest.py`` (lmms-eval is an optional, +undeclared dependency), so these tests always run — in CI without lmms-eval and locally — +exercising the adapter's helpers, guards, and decode loop on a tiny random VLM: no GPU, no +real checkpoint, no network. Real-package fidelity is pinned by the gated contract test in +``tests/integration/``. +""" + +from __future__ import annotations + +import json + +import pytest +import torch +from lmms_eval.api.instance import Instance +from lmms_eval.protocol import ChatMessages +from PIL import Image + +from kempnerforge.config.data import DataConfig +from kempnerforge.config.registry import registry +from kempnerforge.config.schema import AdapterConfig, JobConfig, ModelConfig, VisionEncoderConfig +from kempnerforge.config.video import VideoConfig +from kempnerforge.config.vlm import VLMConfig +from kempnerforge.data.vlm_dataset import ( + DEFAULT_IMAGE_MEAN, + DEFAULT_IMAGE_STD, + pil_to_tensor, +) +from kempnerforge.eval.vlm.adapter import ( + KempnerForgeVLM, + _build_model, + _check_generative, + _ContextBudgetError, + _first_stop, + _frames_to_pixel_values, + _generate_batch, + _load_config, + _load_weights, + _log_checkpoint_metadata, + _render_request, + _resolve_dtype, + _resolve_gen_kwargs, + _to_pil, +) +from kempnerforge.model.vlm import VLMWrapper, build_vlm_wrapper + +DEVICE = torch.device("cpu") + + +# Arch coverage is DERIVED from the registry + the is_generative property (not a +# hardcoded list) so new arches are swept automatically: generative arches get the +# decode/generate sweeps, non-generative ones get the rejection guard. The single +# explicit per-arch truth lives in tests/unit/test_vlm_config.py::TestIsGenerative. +_ALL_VLM_ARCHS = tuple(sorted(registry.list_vlm_configs())) +GENERATIVE_ARCHES = tuple(a for a in _ALL_VLM_ARCHS if VLMConfig.for_arch(a).is_generative) +NON_GENERATIVE_ARCHES = tuple(a for a in _ALL_VLM_ARCHS if not VLMConfig.for_arch(a).is_generative) + +# Per-arch BUILD sizing for a tiny CPU wrapper (sizing only, NOT generativity policy): +# CA needs a cross-attention cadence that fits the tiny layer count. Arches without an +# entry build from defaults; a future arch needing knobs fails the build loudly here. +_ARCH_BUILD_KWARGS = {"cross_attention": {"cross_attention_every_n_layers": 2}} + + +def _vlm_wrapper(arch: str) -> VLMWrapper: + """Build a tiny CPU ``VLMWrapper`` for a generative arch (no checkpoint). + + Uniform ``n_layers=4`` (so CA has at least one cross-attention block) and + ``ffn_hidden_dim=128`` (so MoT's per-modality FFN stays tiny) are valid for every + arch; ``num_image_tokens`` (8) + ``max_text_len`` (32) fits ``max_seq_len`` (64). + """ + mc = ModelConfig( + dim=64, n_layers=4, n_heads=4, vocab_size=256, max_seq_len=64, ffn_hidden_dim=128 + ) + vc = VisionEncoderConfig(type="random", feature_dim=96, num_tokens=8) + lc = VLMConfig.for_arch(arch, max_text_len=32, **_ARCH_BUILD_KWARGS.get(arch, {})) + return build_vlm_wrapper(mc, vc, AdapterConfig(), lc).eval() + + +@pytest.fixture +def arch_wrapper(arch): + """A tiny per-arch ``VLMWrapper``; ``arch`` is provided by ``@pytest.mark.parametrize``.""" + return _vlm_wrapper(arch) + + +class _MockTokenizer: + """Deterministic tokenizer for decode-loop tests (no HF download). + + ``decode`` renders ids as space-joined integers so stop-string and trimming + behavior is easy to assert; ``eos_token_id`` is configurable per test. + """ + + pad_token_id = 0 + + def __init__(self, eos_token_id: int | None = None) -> None: + self.eos_token_id = eos_token_id + + def __call__(self, text: str, add_special_tokens: bool = False) -> dict[str, list[int]]: + del add_special_tokens + return {"input_ids": [(ord(c) % 254) + 1 for c in text]} + + def decode(self, ids: list[int], skip_special_tokens: bool = True) -> str: + del skip_special_tokens + return " ".join(str(int(i)) for i in ids) + + +class _RecordingLogger: + """Captures log calls so metadata/warning behavior is asserted without caplog.""" + + def __init__(self) -> None: + self.infos: list[str] = [] + self.warnings: list[str] = [] + + def info(self, msg: str) -> None: + self.infos.append(msg) + + def warning(self, msg: str) -> None: + self.warnings.append(msg) + + +def _text(s: str) -> dict: + return {"type": "text", "text": s} + + +def _image(img: object) -> dict: + return {"type": "image", "url": img} + + +def _video(url: object) -> dict: + return {"type": "video", "url": url} + + +def _chat(content: list[dict], role: str = "user") -> ChatMessages: + return ChatMessages(messages=[{"role": role, "content": content}]) + + +def _img(size: int = 8) -> Image.Image: + return Image.new("RGB", (size, size), color=(120, 120, 120)) + + +# --------------------------------------------------------------------------- +# _render_request (message -> prompt text + media guards) +# --------------------------------------------------------------------------- + + +class TestRenderRequest: + def test_flattens_text_blocks_in_order(self): + messages = _chat([_text("Question:"), _image(_img()), _text("What color?")]) + images, prompt = _render_request(messages, None) + assert len(images) == 1 + assert prompt == "Question:\nWhat color?" + + def test_video_content_raises(self): + messages = _chat([_text("describe"), {"type": "video", "url": "clip.mp4"}]) + with pytest.raises(NotImplementedError, match="[Vv]ideo"): + _render_request(messages, None) + + def test_audio_content_raises(self): + messages = _chat([_text("listen"), {"type": "audio", "url": "a.wav"}]) + with pytest.raises(NotImplementedError, match="[Aa]udio"): + _render_request(messages, None) + + def test_multi_image_raises(self): + messages = _chat([_text("compare"), _image(_img()), _image(_img())]) + with pytest.raises(NotImplementedError, match="one image"): + _render_request(messages, None) + + def test_no_image_raises(self): + messages = _chat([_text("text only question")]) + with pytest.raises(NotImplementedError, match="one image"): + _render_request(messages, None) + + def test_multi_turn_assistant_raises(self): + messages = ChatMessages( + messages=[ + {"role": "user", "content": [_text("q"), _image(_img())]}, + {"role": "assistant", "content": [_text("a")]}, + {"role": "user", "content": [_text("follow up")]}, + ] + ) + with pytest.raises(NotImplementedError, match="[Mm]ulti-turn"): + _render_request(messages, None) + + +class TestRenderRequestVideo: + """Video-checkpoint rendering (``video_config`` is not None). Decode is stubbed.""" + + @pytest.fixture + def vcfg(self) -> VideoConfig: + return VideoConfig(max_frames=4, min_frames=1, frame_size=16) + + def test_video_decoded_to_frames(self, monkeypatch, vcfg): + frames = [_img(), _img(), _img()] + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.decode_video_frames", lambda path, **kw: frames + ) + out_frames, prompt = _render_request(_chat([_text("describe"), _video("clip.mp4")]), vcfg) + assert out_frames is frames + assert prompt == "describe" + + def test_decode_uses_config_policy(self, monkeypatch, vcfg): + captured: dict = {} + + def _fake(path, **kw): + captured["path"] = path + captured.update(kw) + return [_img()] + + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.decode_video_frames", _fake) + _render_request(_chat([_video("c.mp4")]), vcfg) + assert captured["path"] == "c.mp4" + assert captured["fps"] == vcfg.fps + assert captured["min_frames"] == vcfg.min_frames + assert captured["max_frames"] == vcfg.max_frames + assert captured["sampling_policy"] == vcfg.sampling_policy + + def test_multiple_videos_raise(self, vcfg): + with pytest.raises(NotImplementedError, match="[Mm]ultiple videos"): + _render_request(_chat([_video("a.mp4"), _video("b.mp4")]), vcfg) + + def test_non_path_video_raises(self, vcfg): + msg = _chat([_video({"video": "x", "start": 0.0})]) + with pytest.raises(NotImplementedError, match="path string"): + _render_request(msg, vcfg) + + def test_mixed_image_and_video_raise(self, vcfg): + with pytest.raises(NotImplementedError, match="[Mm]ixed"): + _render_request(_chat([_image(_img()), _video("c.mp4")]), vcfg) + + def test_single_image_is_one_frame_clip(self, vcfg): + frames, prompt = _render_request(_chat([_text("q"), _image(_img())]), vcfg) + assert len(frames) == 1 + assert prompt == "q" + + def test_audio_still_raises(self, vcfg): + msg = _chat([_text("x"), {"type": "audio", "url": "a.wav"}]) + with pytest.raises(NotImplementedError, match="[Aa]udio"): + _render_request(msg, vcfg) + + def test_multi_image_raises(self, vcfg): + with pytest.raises(NotImplementedError, match="exactly one"): + _render_request(_chat([_image(_img()), _image(_img())]), vcfg) + + def test_multi_turn_raises(self, vcfg): + msg = ChatMessages( + messages=[ + {"role": "user", "content": [_text("q"), _video("c.mp4")]}, + {"role": "assistant", "content": [_text("a")]}, + {"role": "user", "content": [_text("again")]}, + ] + ) + with pytest.raises(NotImplementedError, match="[Mm]ulti-turn"): + _render_request(msg, vcfg) + + def test_no_frames_decoded_warns(self, monkeypatch, vcfg): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.decode_video_frames", lambda path, **kw: [] + ) + frames, _ = _render_request(_chat([_video("c.mp4")]), vcfg) + assert frames == [] + assert any("zero clip" in m for m in rec.warnings) + + +# --------------------------------------------------------------------------- +# _frames_to_pixel_values (preprocessing glue) +# --------------------------------------------------------------------------- + + +class TestFramesToPixelValues: + def test_shape_and_dtype(self): + pv = _frames_to_pixel_values([_img()], image_size=16, device=DEVICE, dtype=torch.float32) + assert pv.shape == (1, 3, 16, 16) + assert pv.dtype == torch.float32 + assert pv.device.type == "cpu" + + def test_uses_image_size(self): + pv = _frames_to_pixel_values([_img()], image_size=32, device=DEVICE, dtype=torch.float32) + assert pv.shape == (1, 3, 32, 32) + + def test_parity_with_pil_to_tensor(self): + """The adapter must reuse the exact training preprocessing.""" + img = _img() + pv = _frames_to_pixel_values([img], image_size=24, device=DEVICE, dtype=torch.float32) + expected = pil_to_tensor(img, 24, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + assert torch.allclose(pv[0], expected) + + def test_accepts_path_string(self, tmp_path): + """A frame given as a path string is opened and preprocessed identically.""" + img = _img(16) + path = tmp_path / "frame.png" + img.save(path) + from_path = _frames_to_pixel_values( + [str(path)], image_size=16, device=DEVICE, dtype=torch.float32 + ) + from_pil = _frames_to_pixel_values([img], image_size=16, device=DEVICE, dtype=torch.float32) + assert from_path.shape == (1, 3, 16, 16) + assert torch.allclose(from_path, from_pil) + + +# --------------------------------------------------------------------------- +# _to_pil (str/path -> PIL normalization) + video str-image handling +# --------------------------------------------------------------------------- + + +class TestToPil: + def test_passthrough_pil(self): + img = _img() + assert _to_pil(img) is img + + def test_opens_str_path(self, tmp_path): + from PIL import Image + + path = tmp_path / "frame.png" + _img(16).save(path) + out = _to_pil(str(path)) + assert isinstance(out, Image.Image) + assert out.size == (16, 16) + + +def test_video_checkpoint_accepts_str_image_path( + tmp_path, monkeypatch, tiny_video_configs, tiny_video_vlm_wrapper +): + """A single image supplied as a path string runs on a video checkpoint: it is + opened to PIL before frames_to_clip_tensor (strict PIL-only), so it is processed + rather than silently skipped by the per-request fault handler.""" + _patch_loaders(monkeypatch, _video_job_config(tiny_video_configs), tiny_video_vlm_wrapper) + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32") + path = tmp_path / "frame.png" + _img(16).save(path) + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("describe"), _image(str(path))]}] + + vlm.task_dict = {"t": {"test": {"d0": {}}}} + inst = Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + out = vlm.generate_until([inst]) + assert len(out) == 1 and len(out[0].split()) == 3 # processed, not skipped + + +# --------------------------------------------------------------------------- +# _resolve_gen_kwargs (defaults + task overrides) +# --------------------------------------------------------------------------- + + +class TestResolveGenKwargs: + def test_defaults_are_greedy(self): + r = _resolve_gen_kwargs({}, default_max_new_tokens=64) + assert r == { + "until": [], + "max_new_tokens": 64, + "temperature": 0.0, + "top_k": 0, + "top_p": 1.0, + } + + def test_task_overrides_max_new_tokens(self): + assert _resolve_gen_kwargs({"max_new_tokens": 8}, 64)["max_new_tokens"] == 8 + + def test_sampling_enabled_by_temperature(self): + r = _resolve_gen_kwargs({"temperature": 0.7, "top_p": 0.9, "top_k": 5}, 64) + assert (r["temperature"], r["top_p"], r["top_k"]) == (0.7, 0.9, 5) + + def test_do_sample_false_forces_greedy(self): + r = _resolve_gen_kwargs({"temperature": 0.7, "do_sample": False}, 64) + assert r["temperature"] == 0.0 and r["top_p"] == 1.0 and r["top_k"] == 0 + + def test_until_string_normalized_to_list(self): + assert _resolve_gen_kwargs({"until": "\n\n"}, 64)["until"] == ["\n\n"] + + def test_until_list_preserved(self): + assert _resolve_gen_kwargs({"until": ["a", "b"]}, 64)["until"] == ["a", "b"] + + def test_explicit_zero_max_new_tokens_honored(self): + # 0 is a valid explicit value, not a "missing" fallback. + assert _resolve_gen_kwargs({"max_new_tokens": 0}, 64)["max_new_tokens"] == 0 + + def test_nonpositive_top_p_treated_as_disabled(self): + assert _resolve_gen_kwargs({"temperature": 0.5, "top_p": 0.0}, 64)["top_p"] == 1.0 + assert _resolve_gen_kwargs({"temperature": 0.5, "top_p": -0.3}, 64)["top_p"] == 1.0 + + def test_positive_top_p_honored_when_sampling(self): + assert _resolve_gen_kwargs({"temperature": 0.5, "top_p": 0.1}, 64)["top_p"] == 0.1 + + +# --------------------------------------------------------------------------- +# _generate_batch (cache-less batched decode loop) on a tiny random VLM +# --------------------------------------------------------------------------- + + +def _pixels(batch: int = 1) -> torch.Tensor: + torch.manual_seed(0) + return torch.randn(batch, 3, 16, 16) + + +@pytest.mark.parametrize("arch", GENERATIVE_ARCHES) +class TestGenerateBatchSingle: + """B == 1 must reproduce the v1 single-request behavior (every generative arch).""" + + def _prompt(self) -> list[torch.Tensor]: + return [torch.tensor([5, 9, 12, 3], dtype=torch.long)] + + def test_greedy_is_deterministic(self, arch_wrapper): + pv, pid = _pixels(), self._prompt() + r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) + out1 = _generate_batch(arch_wrapper, _MockTokenizer(), pv, pid, r, 64) + out2 = _generate_batch(arch_wrapper, _MockTokenizer(), pv, pid, r, 64) + assert out1 == out2 and len(out1) == 1 and isinstance(out1[0], str) + + def test_respects_max_new_tokens(self, arch_wrapper): + r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) + out = _generate_batch(arch_wrapper, _MockTokenizer(), _pixels(), self._prompt(), r, 64) + assert len(out[0].split()) == 6 + + def test_until_trims_continuation(self, arch_wrapper): + pv, pid = _pixels(), self._prompt() + one = _generate_batch( + arch_wrapper, + _MockTokenizer(), + pv, + pid, + _resolve_gen_kwargs({"max_new_tokens": 1}, 128), + 64, + )[0] + # decode = space-joined ids, so the first space follows the first token: + # until=[" "] trims to exactly the first generated token. + trimmed = _generate_batch( + arch_wrapper, + _MockTokenizer(), + pv, + pid, + _resolve_gen_kwargs({"max_new_tokens": 6, "until": [" "]}, 128), + 64, + )[0] + assert trimmed == one and " " not in trimmed + + def test_eos_stops_generation(self, arch_wrapper): + pv, pid = _pixels(), self._prompt() + first = _generate_batch( + arch_wrapper, + _MockTokenizer(), + pv, + pid, + _resolve_gen_kwargs({"max_new_tokens": 1}, 128), + 64, + )[0] + out = _generate_batch( + arch_wrapper, + _MockTokenizer(eos_token_id=int(first)), + pv, + pid, + _resolve_gen_kwargs({"max_new_tokens": 6}, 128), + 64, + ) + assert out == [""] + + def test_no_room_raises_context_budget_error(self, arch_wrapper): + # _generate_batch still guards standalone; generate_until converts this into a + # per-task skip (see TestGenerateUntilFaultTolerance). + r = _resolve_gen_kwargs({"max_new_tokens": 100}, 128) + with pytest.raises(_ContextBudgetError, match="max_new_tokens"): + _generate_batch(arch_wrapper, _MockTokenizer(), _pixels(), self._prompt(), r, 64) + + def test_overlong_prompt_is_left_truncated(self, arch_wrapper, monkeypatch): + """A prompt that exceeds the budget (but leaves room) is left-truncated with a warning.""" + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + # Size the prompt from the wrapper's own num_image_tokens (0 for CA, 8 for JD/MoT) + # so it overflows the budget by 6 on every arch — no hardcoded image-token count. + # budget = max_seq_len(64) - num_image_tokens - max_new_tokens(2). + max_new = 2 + budget = 64 - arch_wrapper.num_image_tokens - max_new + long_prompt = [torch.arange(1, budget + 7, dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": max_new}, 128) + out = _generate_batch(arch_wrapper, _MockTokenizer(), _pixels(), long_prompt, r, 64) + assert len(out) == 1 and len(out[0].split()) == 2 + assert any("left-truncating" in m for m in rec.warnings) + + +@pytest.mark.parametrize("arch", GENERATIVE_ARCHES) +class TestGenerateBatchMulti: + """B > 1: right-padding must not change any row's result, and stop + conditions must be tracked per row (every generative arch).""" + + def _prompts(self) -> list[torch.Tensor]: + # Deliberately different lengths so right-padding is exercised. + return [ + torch.tensor([5, 9, 12, 3], dtype=torch.long), + torch.tensor([7, 2], dtype=torch.long), + torch.tensor([1, 4, 8, 11, 20, 6], dtype=torch.long), + ] + + def test_batch_equals_sequential(self, arch_wrapper): + """Key correctness gate: a right-padded batch yields the same per-row + continuation as decoding each request alone (greedy, float32).""" + prompts = self._prompts() + pv = _pixels(len(prompts)) # (3, 3, 16, 16) — one image per request + r = _resolve_gen_kwargs({"max_new_tokens": 5}, 128) + sequential = [ + _generate_batch(arch_wrapper, _MockTokenizer(), pv[i : i + 1], [prompts[i]], r, 64)[0] + for i in range(len(prompts)) + ] + batched = _generate_batch(arch_wrapper, _MockTokenizer(), pv, prompts, r, 64) + assert batched == sequential + + def test_per_row_max_new_tokens(self, arch_wrapper): + prompts = self._prompts() + pv = _pixels(len(prompts)) + r = _resolve_gen_kwargs({"max_new_tokens": 4}, 128) + outs = _generate_batch(arch_wrapper, _MockTokenizer(), pv, prompts, r, 64) + assert len(outs) == 3 and all(len(o.split()) == 4 for o in outs) + + def test_per_row_eos_independent(self, arch_wrapper): + """EOS on one row stops only that row; the batch still returns all rows.""" + prompts = self._prompts() + pv = _pixels(len(prompts)) + first0 = _generate_batch( + arch_wrapper, + _MockTokenizer(), + pv[:1], + [prompts[0]], + _resolve_gen_kwargs({"max_new_tokens": 1}, 128), + 64, + )[0] + outs = _generate_batch( + arch_wrapper, + _MockTokenizer(eos_token_id=int(first0)), + pv, + prompts, + _resolve_gen_kwargs({"max_new_tokens": 5}, 128), + 64, + ) + assert len(outs) == 3 and outs[0] == "" # row 0 stops immediately on EOS + + +# --------------------------------------------------------------------------- +# _generate_batch frame_mask threading (video padded-frame masking) +# --------------------------------------------------------------------------- + + +class _CaptureModel: + """Minimal ``VLMWrapper`` stand-in: records the ``frame_mask`` each forward + receives and returns deterministic zero logits (greedy -> token 0), so the + decode-loop plumbing can be asserted without a real transformer.""" + + def __init__(self, num_image_tokens: int, vocab_size: int = 256) -> None: + self.num_image_tokens = num_image_tokens + self._vocab = vocab_size + self.seen_frame_masks: list[torch.Tensor | None] = [] + + def __call__(self, pixel_values, input_ids, frame_mask=None): + del pixel_values + self.seen_frame_masks.append(frame_mask) + b, t = input_ids.shape + return torch.zeros(b, t, self._vocab), None + + +def test_generate_batch_threads_frame_mask_for_video(): + """A video batch passes the same (B, F) frame_mask into every model forward, so + padded-frame visual tokens are masked from attention exactly as in training.""" + model = _CaptureModel(num_image_tokens=16) + pixel_values = torch.randn(2, 2, 3, 16, 16) # (B, F=2, 3, H, W) + frame_mask = torch.tensor([[True, True], [True, False]]) # row 1: frame 2 padded + prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128) + _generate_batch(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64, frame_mask=frame_mask) + assert model.seen_frame_masks # decode actually ran + assert all(fm is frame_mask for fm in model.seen_frame_masks) + + +def test_generate_batch_no_frame_mask_for_image(): + """An image batch forwards frame_mask=None so the model keeps its unmasked path.""" + model = _CaptureModel(num_image_tokens=8) + pixel_values = torch.randn(1, 3, 16, 16) + prompt_ids = [torch.tensor([5, 9, 12], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 2}, 128) + _generate_batch(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64) + assert model.seen_frame_masks and all(fm is None for fm in model.seen_frame_masks) + + +# --------------------------------------------------------------------------- +# Guards: arch + not-implemented methods +# --------------------------------------------------------------------------- + + +class TestGuards: + @pytest.mark.parametrize("arch", NON_GENERATIVE_ARCHES) + def test_non_generative_arches_rejected(self, arch): + with pytest.raises(ValueError, match="non-causal"): + _check_generative(VLMConfig.for_arch(arch)) + + @pytest.mark.parametrize("arch", GENERATIVE_ARCHES) + def test_generative_arches_allowed(self, arch): + _check_generative(VLMConfig.for_arch(arch)) # must not raise + + def test_loglikelihood_not_implemented(self): + inst = KempnerForgeVLM.__new__(KempnerForgeVLM) # bypass __init__ (no checkpoint needed) + with pytest.raises(NotImplementedError, match="loglikelihood"): + inst.loglikelihood([]) + + def test_multi_round_not_implemented(self): + inst = KempnerForgeVLM.__new__(KempnerForgeVLM) + with pytest.raises(NotImplementedError, match="multi-round"): + inst.generate_until_multi_round([]) + + +# --------------------------------------------------------------------------- +# Loader / __init__ helpers (build a VLM JobConfig + patch the heavy loaders so +# the adapter can be constructed on CPU with no checkpoint). +# --------------------------------------------------------------------------- + + +def _vlm_job_config(tiny_vlm_configs, arch: str | None = None) -> JobConfig: + mc, vc, ac, lc = tiny_vlm_configs + if arch is not None: + # Build the real arch subclass (not a base VLMConfig with a mutated + # .arch) so per-arch config policy like is_generative is exercised. + lc = VLMConfig.for_arch(arch, max_text_len=lc.max_text_len) + return JobConfig( + model=mc, vision_encoder=vc, adapter=ac, vlm=lc, data=DataConfig(tokenizer_path="mock") + ) + + +def _video_job_config(tiny_video_configs) -> JobConfig: + mc, vc, ac, lc, video = tiny_video_configs + return JobConfig( + model=mc, + vision_encoder=vc, + adapter=ac, + vlm=lc, + video=video, + data=DataConfig(tokenizer_path="mock"), + ) + + +def _patch_loaders(monkeypatch, job: JobConfig, model) -> None: + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._load_config", lambda _p: job) + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._load_weights", lambda *a, **k: model) + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.build_tokenizer", lambda _p: _MockTokenizer() + ) + + +# --------------------------------------------------------------------------- +# _resolve_dtype +# --------------------------------------------------------------------------- + + +class TestResolveDtype: + @pytest.mark.parametrize( + "name,expected", + [("bfloat16", torch.bfloat16), ("float16", torch.float16), ("float32", torch.float32)], + ) + def test_string_dtypes(self, name, expected): + assert _resolve_dtype(name) == expected + + def test_passthrough_torch_dtype(self): + assert _resolve_dtype(torch.float64) == torch.float64 + + def test_unsupported_dtype_raises(self): + with pytest.raises(ValueError, match="Unsupported dtype"): + _resolve_dtype("float64") + + +# --------------------------------------------------------------------------- +# _first_stop +# --------------------------------------------------------------------------- + + +class TestFirstStop: + def test_no_match_returns_none(self): + assert _first_stop("hello world", ["xyz"]) is None + + def test_empty_until_returns_none(self): + assert _first_stop("abc", []) is None + + def test_earliest_of_multiple(self): + # "cd" at index 2 precedes "ef" at index 4. + assert _first_stop("abcdef", ["ef", "cd"]) == 2 + + def test_single_match_index(self): + assert _first_stop("a.b", ["."]) == 1 + + +# --------------------------------------------------------------------------- +# _log_checkpoint_metadata +# --------------------------------------------------------------------------- + + +class TestLogCheckpointMetadata: + def test_missing_metadata_is_noop(self, tmp_path, monkeypatch): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + _log_checkpoint_metadata(tmp_path) + assert rec.infos == [] and rec.warnings == [] + + def test_valid_metadata_logged(self, tmp_path, monkeypatch): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + (tmp_path / "metadata.json").write_text(json.dumps({"step": 7, "tokens_seen": 1234})) + _log_checkpoint_metadata(tmp_path) + assert any("step=7" in m and "tokens_seen=1234" in m for m in rec.infos) + + def test_malformed_metadata_warns_not_raises(self, tmp_path, monkeypatch): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + (tmp_path / "metadata.json").write_text("{not valid json") + _log_checkpoint_metadata(tmp_path) # must not raise + assert any("Could not read" in m for m in rec.warnings) + + +# --------------------------------------------------------------------------- +# _load_config (VLM-only guard) +# --------------------------------------------------------------------------- + + +class TestLoadConfig: + def test_non_vlm_config_rejected(self, monkeypatch, tiny_job_config): + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.load_config", + lambda _p, cli_args=None: tiny_job_config, + ) + with pytest.raises(ValueError, match="not a VLM config"): + _load_config("ignored.toml") + + def test_vlm_config_accepted(self, monkeypatch, tiny_vlm_configs): + job = _vlm_job_config(tiny_vlm_configs) + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.load_config", lambda _p, cli_args=None: job + ) + assert _load_config("ignored.toml") is job + + +# --------------------------------------------------------------------------- +# _load_weights (path resolution / missing checkpoint) +# --------------------------------------------------------------------------- + + +class TestLoadWeights: + def test_missing_checkpoint_raises(self, tmp_path, tiny_vlm_configs): + job = _vlm_job_config(tiny_vlm_configs) + missing = tmp_path / "no_such_checkpoint" + with pytest.raises(FileNotFoundError, match="does not exist"): + _load_weights(job, str(missing), torch.device("cpu"), torch.float32) + + +# --------------------------------------------------------------------------- +# KempnerForgeVLM.__init__ guards +# --------------------------------------------------------------------------- + + +class TestInitGuards: + @pytest.mark.parametrize("arch", NON_GENERATIVE_ARCHES) + def test_non_generative_fails_fast_before_load(self, monkeypatch, tiny_vlm_configs, arch): + job = _vlm_job_config(tiny_vlm_configs, arch=arch) + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._load_config", lambda _p: job) + + def _must_not_load(*args, **kwargs): + raise AssertionError("model load must not run for an unsupported arch") + + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._load_weights", _must_not_load) + with pytest.raises(ValueError, match="non-causal"): + KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32") + + def test_ignored_kwargs_warn(self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + _patch_loaders(monkeypatch, _vlm_job_config(tiny_vlm_configs), tiny_vlm_wrapper) + KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32", bogus=1) + assert any("bogus" in m for m in rec.warnings) + + @pytest.mark.parametrize("arch", GENERATIVE_ARCHES) + def test_init_populates_attrs(self, monkeypatch, tiny_vlm_configs, arch): + job = _vlm_job_config(tiny_vlm_configs, arch=arch) + _patch_loaders(monkeypatch, job, _vlm_wrapper(arch)) + vlm = KempnerForgeVLM( + config="x", checkpoint="y", device="cpu", dtype="float32", batch_size=2 + ) + assert vlm._arch == arch + assert vlm._max_seq_len == job.model.max_seq_len + assert vlm._batch_size == 2 + assert vlm._dtype == torch.float32 + assert vlm._is_video is False + assert vlm._frames_per_clip == 1 + assert vlm._frame_size == job.data.hf_image_size + + def test_dtype_defaults_from_config(self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper): + # No explicit dtype -> use the checkpoint config's train.param_dtype. + from kempnerforge.config.training import TrainConfig + + mc, vc, ac, lc = tiny_vlm_configs + job = JobConfig( + model=mc, + vision_encoder=vc, + adapter=ac, + vlm=lc, + data=DataConfig(tokenizer_path="mock"), + train=TrainConfig(mixed_precision="fp32"), + ) + _patch_loaders(monkeypatch, job, tiny_vlm_wrapper) + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu") + assert vlm._dtype == job.train.param_dtype == torch.float32 + + def test_explicit_dtype_overrides_config(self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper): + # An explicit dtype wins over the config default. + _patch_loaders(monkeypatch, _vlm_job_config(tiny_vlm_configs), tiny_vlm_wrapper) + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float16") + assert vlm._dtype == torch.float16 + + +# --------------------------------------------------------------------------- +# _build_model (frames_per_clip wiring for image vs video checkpoints) +# --------------------------------------------------------------------------- + + +class TestBuildModel: + def test_video_config_sets_frames_per_clip(self, tiny_video_configs): + job = _video_job_config(tiny_video_configs) + model = _build_model(job, torch.device("cpu"), torch.float32) + assert job.video is not None + assert model.frames_per_clip == job.video.max_frames == 2 + + def test_image_config_frames_per_clip_is_one(self, tiny_vlm_configs): + job = _vlm_job_config(tiny_vlm_configs) + model = _build_model(job, torch.device("cpu"), torch.float32) + assert model.frames_per_clip == 1 + + +# --------------------------------------------------------------------------- +# generate_until (end-to-end on a tiny VLM via the fake lmms-eval objects) +# --------------------------------------------------------------------------- + + +class TestGenerateUntil: + @pytest.mark.parametrize("arch", GENERATIVE_ARCHES) + def test_batches_and_restores_order(self, monkeypatch, tiny_vlm_configs, arch): + _patch_loaders( + monkeypatch, _vlm_job_config(tiny_vlm_configs, arch=arch), _vlm_wrapper(arch) + ) + vlm = KempnerForgeVLM( + config="x", checkpoint="y", device="cpu", dtype="float32", batch_size=2 + ) + + img = _img() + + def doc_to_messages(doc): + return [{"role": "user", "content": [_text(doc["q"]), _image(img)]}] + + vlm.task_dict = { + "t": {"test": {"d0": {"q": "hi?"}, "d1": {"q": "describe the picture please"}}} + } + # Different context strings (args[0]) so the Collator reorders the batch by length; + # get_original must then restore the original request order. + specs = [("d0", "c"), ("d1", "cc")] + instances = [ + Instance( + request_type="generate_until", + arguments=(ctx, doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, (doc_id, ctx) in enumerate(specs) + ] + + # Each request decoded alone (greedy → deterministic), in original order. + singles = [vlm.generate_until([inst])[0] for inst in instances] + batched = vlm.generate_until(instances) + + assert batched == singles # batching + reorder + get_original preserve per-request results + assert all(isinstance(o, str) for o in batched) + assert all(len(o.split()) == 3 for o in batched) # greedy emits exactly max_new_tokens + + +class TestGenerateUntilVideo: + """End-to-end video decode: stack -> (B, F, 3, H, W) -> forward -> strings.""" + + def test_video_batches_and_restores_order( + self, monkeypatch, tiny_video_configs, tiny_video_vlm_wrapper + ): + job = _video_job_config(tiny_video_configs) + _patch_loaders(monkeypatch, job, tiny_video_vlm_wrapper) + # Two real frames per clip; zero-padded to frames_per_clip (== 2) downstream. + monkeypatch.setattr( + "kempnerforge.eval.vlm.adapter.decode_video_frames", + lambda path, **kw: [_img(), _img()], + ) + vlm = KempnerForgeVLM( + config="x", checkpoint="y", device="cpu", dtype="float32", batch_size=2 + ) + assert vlm._is_video is True + assert vlm._frames_per_clip == 2 + + def doc_to_messages(doc): + return [{"role": "user", "content": [_text(doc["q"]), _video(doc["v"])]}] + + vlm.task_dict = { + "t": { + "test": { + "d0": {"q": "what?", "v": "a.mp4"}, + "d1": {"q": "describe it", "v": "b.mp4"}, + } + } + } + specs = [("d0", "c"), ("d1", "cc")] + instances = [ + Instance( + request_type="generate_until", + arguments=(ctx, doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, (doc_id, ctx) in enumerate(specs) + ] + # Decoding each request alone must match the batched (stacked, 5-D) result; if the + # adapter folded frames with cat, the forward would trip the frames-per-clip check. + singles = [vlm.generate_until([inst])[0] for inst in instances] + batched = vlm.generate_until(instances) + assert batched == singles + assert all(isinstance(o, str) and len(o.split()) == 3 for o in batched) + + def test_video_pixel_values_assembled_5d(self): + """The video batch is (B, F, 3, H, W) via stack, not (B*F, 3, H, W).""" + from kempnerforge.data.vlm_dataset import frames_to_clip_tensor + + clip0, _ = frames_to_clip_tensor([_img(), _img()], max_frames=2, frame_size=16) + clip1, _ = frames_to_clip_tensor([_img()], max_frames=2, frame_size=16) + pixel_values = torch.stack([clip0, clip1], dim=0) + assert pixel_values.shape == (2, 2, 3, 16, 16) + assert pixel_values.ndim == 5 + + +# --------------------------------------------------------------------------- +# generate_until fault tolerance + empty-prompt guard +# --------------------------------------------------------------------------- + + +class TestGenerateUntilFaultTolerance: + """A request that fails to render/preprocess is isolated (empty output + a + warning) so the rest of the batch still scores; an empty flattened prompt is + guarded rather than crashing the decode. The per-request handler catches any + ``Exception``, so an unexpected error skips just that doc; a fatal + ``BaseException`` (Ctrl-C / ``SystemExit``) still propagates so the run stays + interruptible rather than being silently scored "".""" + + def _vlm(self, monkeypatch, tiny_vlm_configs, batch_size=2): + _patch_loaders( + monkeypatch, _vlm_job_config(tiny_vlm_configs), _vlm_wrapper("joint_decoder") + ) + return KempnerForgeVLM( + config="x", checkpoint="y", device="cpu", dtype="float32", batch_size=batch_size + ) + + def test_bad_request_skipped_others_complete(self, monkeypatch, tiny_vlm_configs): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + vlm = self._vlm(monkeypatch, tiny_vlm_configs) + img = _img() + + def doc_to_messages(doc): + # The "bad" doc carries two images (multi-image -> NotImplementedError in + # _render_request); the "good" doc carries exactly one. + content = [_text(doc["q"])] + [_image(img)] * doc["n_images"] + return [{"role": "user", "content": content}] + + vlm.task_dict = { + "t": { + "test": { + "bad": {"q": "compare", "n_images": 2}, + "good": {"q": "describe", "n_images": 1}, + } + } + } + instances = [ + Instance( + request_type="generate_until", + arguments=(ctx, doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, (doc_id, ctx) in enumerate([("bad", "c"), ("good", "cc")]) + ] + outs = vlm.generate_until(instances) + assert outs[0] == "" # bad request isolated (order restored by get_original) + assert len(outs[1].split()) == 3 # good request completes normally + assert any("Skipping request" in m and "doc_id=bad" in m for m in rec.warnings) + + def test_empty_prompt_is_guarded(self, monkeypatch, tiny_vlm_configs): + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + vlm = self._vlm(monkeypatch, tiny_vlm_configs, batch_size=1) + + def doc_to_messages(doc): + del doc + # One image, no text block -> empty flattened prompt. + return [{"role": "user", "content": [_image(_img())]}] + + vlm.task_dict = {"t": {"test": {"d0": {}}}} + inst = Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + assert vlm.generate_until([inst]) == [""] + assert any("empty prompt" in m for m in rec.warnings) + + def test_all_requests_bad_returns_empties(self, monkeypatch, tiny_vlm_configs): + # Every request in the chunk fails -> generation is short-circuited (no empty + # stack/cat) and each slot returns "", order preserved. + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", _RecordingLogger()) + vlm = self._vlm(monkeypatch, tiny_vlm_configs) + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("x"), _image(_img()), _image(_img())]}] + + vlm.task_dict = {"t": {"test": {"d0": {}, "d1": {}}}} + instances = [ + Instance( + request_type="generate_until", + arguments=(ctx, doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, (doc_id, ctx) in enumerate([("d0", "c"), ("d1", "cc")]) + ] + assert vlm.generate_until(instances) == ["", ""] + + def test_unexpected_error_isolated(self, monkeypatch, tiny_vlm_configs): + # The per-request handler catches Exception (broadened from the original + # NotImplementedError/OSError/ValueError), so an unexpected error on one + # document — here a RuntimeError from _render_request — is isolated as "" + # with a warning rather than aborting the whole run. + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + vlm = self._vlm(monkeypatch, tiny_vlm_configs, batch_size=1) + + def boom(*_a, **_k): + raise RuntimeError("boom") + + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._render_request", boom) + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("hi"), _image(_img())]}] + + vlm.task_dict = {"t": {"test": {"d0": {}}}} + inst = Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + assert vlm.generate_until([inst]) == [""] + assert any("Skipping request" in m and "doc_id=d0" in m for m in rec.warnings) + + def test_fatal_baseexception_propagates(self, monkeypatch, tiny_vlm_configs): + # The broadened catch is `except Exception`, so a BaseException (a Ctrl-C + # KeyboardInterrupt or a SystemExit) is NOT swallowed: it surfaces rather + # than being silently scored "", keeping a long eval run interruptible. + vlm = self._vlm(monkeypatch, tiny_vlm_configs, batch_size=1) + + def interrupt(*_a, **_k): + raise KeyboardInterrupt + + monkeypatch.setattr("kempnerforge.eval.vlm.adapter._render_request", interrupt) + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("hi"), _image(_img())]}] + + vlm.task_dict = {"t": {"test": {"d0": {}}}} + inst = Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + with pytest.raises(KeyboardInterrupt): + vlm.generate_until([inst]) + + def test_missing_image_path_skipped(self, monkeypatch, tiny_vlm_configs): + # A bad image path raises FileNotFoundError (subclass of OSError) inside _to_pil; + # it is isolated as "" like other per-document failures. + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + vlm = self._vlm(monkeypatch, tiny_vlm_configs, batch_size=1) + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("describe"), _image("/no/such/frame.png")]}] + + vlm.task_dict = {"t": {"test": {"d0": {}}}} + inst = Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + assert vlm.generate_until([inst]) == [""] + assert any("Skipping request" in m and "doc_id=d0" in m for m in rec.warnings) + + def test_over_budget_chunk_skipped_not_aborted(self, monkeypatch, tiny_vlm_configs): + # A task whose max_new_tokens over-budgets the context skips its own requests + # (via _ContextBudgetError) instead of aborting the whole run. + rec = _RecordingLogger() + monkeypatch.setattr("kempnerforge.eval.vlm.adapter.logger", rec) + vlm = self._vlm(monkeypatch, tiny_vlm_configs) # tiny max_seq_len == 64 + img = _img() + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("hi"), _image(img)]}] + + vlm.task_dict = {"t": {"test": {"d0": {}, "d1": {}}}} + instances = [ + Instance( + request_type="generate_until", + arguments=(ctx, doc_to_messages, {"max_new_tokens": 1000}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, (doc_id, ctx) in enumerate([("d0", "c"), ("d1", "cc")]) + ] + assert vlm.generate_until(instances) == ["", ""] + assert any("Skipping" in m and "max_new_tokens" in m for m in rec.warnings) diff --git a/tests/unit/eval/vlm/test_import_isolation.py b/tests/unit/eval/vlm/test_import_isolation.py new file mode 100644 index 0000000..feec89f --- /dev/null +++ b/tests/unit/eval/vlm/test_import_isolation.py @@ -0,0 +1,25 @@ +"""Import isolation: ``import kempnerforge`` must not require lmms-eval. + +The eval subpackage depends on the optional, undeclared ``lmms-eval`` package, +so it must never be imported on the default import path. This test runs in a +fresh subprocess (so prior imports in the pytest session do not pollute +``sys.modules``) and asserts that importing the main package — and even the +eval namespace packages — does not pull in ``lmms_eval``. It intentionally does +NOT skip when lmms-eval is installed: the property must hold either way. +""" + +from __future__ import annotations + +import subprocess +import sys + + +def test_import_kempnerforge_does_not_import_lmms_eval(): + code = ( + "import sys, kempnerforge, kempnerforge.eval, kempnerforge.eval.vlm; " + "assert 'lmms_eval' not in sys.modules, 'importing kempnerforge pulled in lmms_eval'; " + "print('ISOLATED')" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + assert "ISOLATED" in result.stdout diff --git a/tests/unit/eval/vlm/test_manifest.py b/tests/unit/eval/vlm/test_manifest.py new file mode 100644 index 0000000..4e59c3c --- /dev/null +++ b/tests/unit/eval/vlm/test_manifest.py @@ -0,0 +1,20 @@ +"""CPU unit test for the lmms-eval registration manifest. + +Runs against the fake ``lmms_eval`` injected by ``conftest.py`` (lmms-eval is an optional, +undeclared dependency), so it executes in CI and covers ``manifest.py``. Real entry-point +resolution against the installed package is verified by the gated integration test. +""" + +from __future__ import annotations + +from lmms_eval.models.registry_v2 import ModelManifest + +from kempnerforge.eval.vlm.manifest import MANIFEST + + +def test_manifest_is_well_formed(): + assert isinstance(MANIFEST, ModelManifest) + assert MANIFEST.model_id == "kempnerforge_vlm" + assert MANIFEST.chat_class_path == "kempnerforge.eval.vlm.adapter.KempnerForgeVLM" + # Chat-only registration keeps resolution correct even under force_simple. + assert MANIFEST.simple_class_path is None diff --git a/tests/unit/test_video_dataset.py b/tests/unit/test_video_dataset.py index 45f45ef..1085b1d 100644 --- a/tests/unit/test_video_dataset.py +++ b/tests/unit/test_video_dataset.py @@ -161,6 +161,26 @@ def test_len(self): ds = _StubVideoDataset(["1", "2", "3"], ["a", "b", "c"]) assert len(ds) == 3 + def test_pixel_values_match_shared_helper(self, monkeypatch): + """The dataset's clip is exactly ``frames_to_clip_tensor``'s output (the + shared helper), guaranteeing training/eval frame-packing parity by + construction.""" + from kempnerforge.data.vlm_dataset import frames_to_clip_tensor + + frames = _frames(3) + monkeypatch.setattr(vd, "decode_video_frames", lambda *a, **k: frames) + ds = _StubVideoDataset(["1"], ["a cat."], max_frames=8, frame_size=16) + item = ds[0] + expected_pv, expected_mask = frames_to_clip_tensor( + frames, + max_frames=8, + frame_size=16, + image_mean=DEFAULT_IMAGE_MEAN, + image_std=DEFAULT_IMAGE_STD, + ) + assert torch.equal(item["pixel_values"], expected_pv) + assert torch.equal(item["frame_mask"], expected_mask) + # --------------------------------------------------------------------------- # VideoCollator diff --git a/tests/unit/test_vlm_config.py b/tests/unit/test_vlm_config.py index 5e865eb..82ecb18 100644 --- a/tests/unit/test_vlm_config.py +++ b/tests/unit/test_vlm_config.py @@ -119,6 +119,22 @@ def test_registry_has_all_archs(self): assert {"joint_decoder", "cross_attention", "mot", "moma"} <= archs +class TestIsGenerative: + """Per-arch ``is_generative`` capability, queried by generation-only + consumers (the lmms-eval chat adapter). Generative arches inherit the base + ``True``; MoMa overrides to ``False`` (non-causal expert-choice routing).""" + + @pytest.mark.parametrize( + "config", + [VLMConfig(), JointDecoderConfig(), CrossAttentionConfig(), MoTConfig()], + ) + def test_generative_arches_are_generative(self, config): + assert config.is_generative is True + + def test_moma_is_not_generative(self): + assert MoMaConfig().is_generative is False + + class TestCrossAttentionResolvedHeads: """`cross_attention_n_heads=0` and `cross_attention_n_kv_heads=0` both resolve against ModelConfig.n_heads at build time. The block diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index b9ee4dc..d714050 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -11,8 +11,9 @@ DEFAULT_IMAGE_STD, HuggingFaceVLMDataset, VLMCollator, - _pil_to_tensor, _tokenize_and_mask, + frames_to_clip_tensor, + pil_to_tensor, ) @@ -21,38 +22,38 @@ def _make_image(size: int = 64, color: int = 128) -> Image.Image: # --------------------------------------------------------------------------- -# _pil_to_tensor +# pil_to_tensor # --------------------------------------------------------------------------- class TestPILToTensor: def test_shape_and_dtype(self): img = _make_image(48) - t = _pil_to_tensor(img, 224, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + t = pil_to_tensor(img, 224, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) assert t.shape == (3, 224, 224) assert t.dtype == torch.float32 def test_non_square_resize(self): img = Image.new("RGB", (100, 50), color=(0, 128, 255)) - t = _pil_to_tensor(img, 64, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + t = pil_to_tensor(img, 64, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) assert t.shape == (3, 64, 64) def test_grayscale_promoted_to_rgb(self): img = Image.new("L", (32, 32), color=128) - t = _pil_to_tensor(img, 32, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + t = pil_to_tensor(img, 32, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) assert t.shape == (3, 32, 32) def test_normalization(self): """Input with pixels at 127.5 (0.5 after /255) should normalize to ~0 under mean=0.5, std=0.5.""" img = Image.new("RGB", (16, 16), color=(128, 128, 128)) # ~0.502 - t = _pil_to_tensor(img, 16, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + t = pil_to_tensor(img, 16, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) # (0.502 - 0.5) / 0.5 ~= 0.004 assert abs(float(t.mean())) < 0.02 def test_non_pil_raises(self): with pytest.raises(TypeError, match="Expected a PIL.Image"): - _pil_to_tensor(torch.randn(3, 16, 16), 16, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + pil_to_tensor(torch.randn(3, 16, 16), 16, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) # --------------------------------------------------------------------------- @@ -443,3 +444,43 @@ def __init__(self, ds, tokenizer_path, max_text_len): assert batch["input_ids"].shape == (2, 12) assert batch["labels"].shape == (2, 12) assert batch["pixel_values"].shape == (2, 3, 16, 16) + + +class TestFramesToClipTensor: + """The shared frame-packing helper used by both training (``WebVidVideoDataset``) + and the VLM evaluation adapter.""" + + def test_full_clip_shape_and_dtype(self): + frames = [_make_image(32) for _ in range(4)] + pv, mask = frames_to_clip_tensor(frames, max_frames=4, frame_size=16) + assert pv.shape == (4, 3, 16, 16) + assert pv.dtype == torch.float32 + assert mask.tolist() == [True, True, True, True] + + def test_short_clip_is_zero_padded(self): + frames = [_make_image(32) for _ in range(2)] + pv, mask = frames_to_clip_tensor(frames, max_frames=5, frame_size=16) + assert pv.shape == (5, 3, 16, 16) + assert mask.tolist() == [True, True, False, False, False] + assert torch.count_nonzero(pv[2:]) == 0 # padded frames are zero + + def test_single_image_is_one_frame_clip(self): + pv, mask = frames_to_clip_tensor([_make_image(20)], max_frames=3, frame_size=16) + assert pv.shape == (3, 3, 16, 16) + assert mask.tolist() == [True, False, False] + + def test_per_frame_matches_pil_to_tensor(self): + img = _make_image(40) + pv, _ = frames_to_clip_tensor([img], max_frames=1, frame_size=16) + expected = pil_to_tensor(img, 16, DEFAULT_IMAGE_MEAN, DEFAULT_IMAGE_STD) + assert torch.allclose(pv[0], expected) + + def test_uses_frame_size(self): + pv, _ = frames_to_clip_tensor([_make_image(32)], max_frames=2, frame_size=24) + assert pv.shape == (2, 3, 24, 24) + + def test_excess_frames_truncated_to_budget(self): + frames = [_make_image(16) for _ in range(5)] + pv, mask = frames_to_clip_tensor(frames, max_frames=3, frame_size=16) + assert pv.shape == (3, 3, 16, 16) + assert mask.all()