diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 46cd90a..7a85d83 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -60,6 +60,15 @@ result. the original wearer probabilities, excludes wearer backchannels from headline ENTRY FPR, and records unanchored/ambiguous cases. +### Implemented but not GPU-verified + +- MiniCPM-o 4.5 input-only feature extraction has two unmeasured modes: a + native 1 Hz state causally held on the 5 Hz grid, and a 5 Hz completed-cache + plus temporary-partial-tail branch. The latter refuses the normal + head-loading contract unless boundary, replay, snapshot-restoration, and + future-audio checks pass. A real-GPU smoke run is still required; see + [`docs/minicpmo_streaming_features.md`](docs/minicpmo_streaming_features.md). + ### Measured screening result - The released EgoSpeak MiniROADA A+V checkpoint on Sessions 1--3 obtains diff --git a/README.md b/README.md index 48d0c3c..b7de0ea 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,7 @@ change what the Stage-0 result measures. | `easycom_manifest.py` | Verified on EasyCom | Discover one-minute clips, preserve session timecodes, and check 5 Hz targets | | `qwen_omni_features.py` | Implemented, not executed here | Load only the 3B Thinker and extract fused states from one AV window | | `streaming_turn_model.py` | Verified on EasyCom | Produce one causal Omni state every 200 ms, with causal-prefix and open-stream modes | +| `streaming_minicpmo_features_directly.py` | Implemented; real-GPU smoke pending | Extract input-only MiniCPM-o 4.5 native 1 Hz held states or audited 5 Hz cache-plus-partial-tail states | | `TurnPredictionHead` | Trained and evaluated | Map `[2048]` to `[10, 3]` logits | | `train_turn_head.py` | Verified on EasyCom | Train the frozen-Omni anticipation head with the official EasyCom split | | `infer_turn_head.py` | Verified on EasyCom | Save EgoSpeak-format predictions and reproduce per-horizon AP/mAP | diff --git a/docs/minicpmo_streaming_features.md b/docs/minicpmo_streaming_features.md new file mode 100644 index 0000000..9379af5 --- /dev/null +++ b/docs/minicpmo_streaming_features.md @@ -0,0 +1,204 @@ +# MiniCPM-o 4.5 causal global features + +`streaming_minicpmo_features_directly.py` uses MiniCPM-o 4.5 only as a frozen +audio--visual feature extractor. It never calls `streaming_generate`, never +loads TTS, and never lets model-generated speech/listen tokens enter the +conversation history. + +## Frozen modes + +### `native-1hz-hold` + +- The model commits one input-only AV chunk per completed second. +- The newest completed state is forward-held onto the 5 Hz target grid. +- The first four rows (0.2--0.8 s) use an all-zero initialization. The 1.0-s + state is never backfilled into those rows. +- `state_source_times`, `native_update_mask`, and `initialization_mask` make + this alignment auditable. + +### `cached-partial-5hz` + +- Completed one-second chunks remain in the official MiniCPM streaming cache. +- At a non-boundary 200-ms cutoff, the extractor snapshots the LLM cache, + audio-encoder cache, streaming Mel buffer, chunk index, and RNG state through + MiniCPM's official speculative-state API. +- It appends only the current half-open AV tail, reads the final normalized + Qwen3 hidden state, and restores the snapshot. Temporary closure never enters + the main cache. +- Native commits retain MiniCPM's normal streaming-window policy. Destructive + round pruning is suspended only inside the temporary transaction because the + upstream snapshot can truncate appended KV tokens but cannot recreate a + previously pruned round. +- At a one-second boundary, a complete native chunk is committed. + +The 5 Hz branch is saved as the normal `mode=causal-prefix` contract only after +all mandatory checks pass: + +1. a temporary full chunk matches the native committed state at a one-second + boundary; +2. a clean stream replay matches the cached result; +3. changing audio strictly after the cutoff leaves the current state exactly + unchanged; +4. snapshot restoration leaves LLM/audio cache lengths and checksums, chunk + history, position offset, and the audio chunk index unchanged. + +With `--allow-unverified-partial-tail`, a failed or too-short audit may be +saved only as `mode=minicpmo-unverified-partial-tail`; the default head loader +rejects that artifact. + +## Environment + +The inspected upstream revision requires Python 3.10 and +`transformers==4.51.0`. Use a dedicated environment rather than changing the +working Qwen environment. + +```bash +python -m pip install \ + "transformers==4.51.0" \ + accelerate \ + "torch>=2.3.0,<=2.8.0" \ + "torchaudio<=2.8.0" \ + "minicpmo-utils>=1.0.5" \ + "qwen-omni-utils[decord]" +``` + +The commands below use the upstream revision inspected during implementation, +`1f761131fa83f5ed3cd6f2f22b225c4501d154fa`. If the model is already stored +locally, replace `--model-id` with that snapshot directory and record its +immutable revision separately. + +## Required smoke sequence + +First validate the source timeline and decoder without loading MiniCPM: + +```bash +python streaming_minicpmo_features_directly.py \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_4 \ + --end 3 \ + --mode native-1hz-hold \ + --decoder decord \ + --media-only +``` + +Then extract three seconds with the native held-state baseline: + +```bash +CUDA_VISIBLE_DEVICES=0 python streaming_minicpmo_features_directly.py \ + --model-id openbmb/MiniCPM-o-4_5 \ + --revision 1f761131fa83f5ed3cd6f2f22b225c4501d154fa \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_4 \ + --end 3 \ + --mode native-1hz-hold \ + --dtype bfloat16 \ + --attention sdpa \ + --device cuda:0 \ + --decoder decord \ + --decoder-threads 1 \ + --max-open-readers 1 \ + --output /data1/yfzx/wangx/Datasets/easycom/minicpmo_features/smoke/Session_4_native_3s.pt +``` + +Run the verified cached partial-tail smoke test separately: + +```bash +CUDA_VISIBLE_DEVICES=0 python streaming_minicpmo_features_directly.py \ + --model-id openbmb/MiniCPM-o-4_5 \ + --revision 1f761131fa83f5ed3cd6f2f22b225c4501d154fa \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_4 \ + --end 3 \ + --mode cached-partial-5hz \ + --verify-every 1 \ + --future-intervention-every 1 \ + --dtype bfloat16 \ + --attention sdpa \ + --device cuda:0 \ + --decoder decord \ + --decoder-threads 1 \ + --max-open-readers 1 \ + --output /data1/yfzx/wangx/Datasets/easycom/minicpmo_features/smoke/Session_4_partial_3s.pt +``` + +Do not add `--allow-unverified-partial-tail` to a paper run. A verification +failure is a protocol failure to diagnose, not a tolerance to bypass. + +Inspect both artifacts before a full session: + +```bash +python - <<'PY' +from pathlib import Path +import torch + +root = Path("/data1/yfzx/wangx/Datasets/easycom/minicpmo_features/smoke") +for name in ("Session_4_native_3s.pt", "Session_4_partial_3s.pt"): + payload = torch.load(root / name, map_location="cpu", weights_only=False) + print(name) + print(" mode:", payload["mode"], payload["extractor_mode"]) + print(" states:", tuple(payload["states"].shape)) + print(" times:", payload["times"][:6].tolist()) + print(" source times:", payload["state_source_times"][:6].tolist()) + print(" verified:", payload["partial_tail_admissible"]) + print(" checks:", payload["verification"]) +PY +``` + +Expected alignment: + +```text +native times: 0.2 0.4 0.6 0.8 1.0 1.2 ... +native source time: 0.0 0.0 0.0 0.0 1.0 1.0 ... +partial source time: 0.2 0.4 0.6 0.8 1.0 1.2 ... +``` + +## Full Session 4 + +Native 1 Hz held features: + +```bash +mkdir -p /data1/yfzx/wangx/Datasets/easycom/minicpmo_features/native_1hz_hold + +CUDA_VISIBLE_DEVICES=0 python streaming_minicpmo_features_directly.py \ + --model-id openbmb/MiniCPM-o-4_5 \ + --revision 1f761131fa83f5ed3cd6f2f22b225c4501d154fa \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_4 \ + --target /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame/Session_4.npy \ + --mode native-1hz-hold \ + --dtype bfloat16 \ + --attention sdpa \ + --device cuda:0 \ + --decoder decord \ + --decoder-threads 1 \ + --max-open-readers 1 \ + --frame-cache-size 32 \ + --output /data1/yfzx/wangx/Datasets/easycom/minicpmo_features/native_1hz_hold/Session_4.pt +``` + +Verified 5 Hz cached partial-tail features: + +```bash +mkdir -p /data1/yfzx/wangx/Datasets/easycom/minicpmo_features/cached_partial_5hz + +CUDA_VISIBLE_DEVICES=0 python streaming_minicpmo_features_directly.py \ + --model-id openbmb/MiniCPM-o-4_5 \ + --revision 1f761131fa83f5ed3cd6f2f22b225c4501d154fa \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_4 \ + --target /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame/Session_4.npy \ + --mode cached-partial-5hz \ + --dtype bfloat16 \ + --attention sdpa \ + --device cuda:0 \ + --decoder decord \ + --decoder-threads 1 \ + --max-open-readers 1 \ + --frame-cache-size 32 \ + --output /data1/yfzx/wangx/Datasets/easycom/minicpmo_features/cached_partial_5hz/Session_4.pt +``` + +The full run still executes the first clean-replay and future-intervention +checks plus every one-second boundary check. Add `--verify-every 250` and +`--future-intervention-every 250` for periodic 50-second audits; smaller values +make long sessions increasingly expensive because a clean audit replays the +prefix from the session start. + +After extraction, either directory can be passed directly as +`train_turn_head.py --feature-root ...` because verified artifacts retain the +existing dense global-feature contract. diff --git a/streaming_minicpmo_features_directly.py b/streaming_minicpmo_features_directly.py new file mode 100644 index 0000000..0c65f00 --- /dev/null +++ b/streaming_minicpmo_features_directly.py @@ -0,0 +1,1089 @@ +#!/usr/bin/env python3 +"""Extract strictly causal MiniCPM-o 4.5 states from original AV media. + +MiniCPM-o is used only as a frozen input feature extractor. This script never +calls ``streaming_generate`` and therefore never inserts model-generated +listen/speak/text/audio tokens into the recorded conversation. + +Two representations are implemented: + +``native-1hz-hold`` + Commit one native one-second AV chunk at each integer-second boundary and + forward-hold the newest completed state on the 5 Hz label grid. Rows + before the first completed chunk use a frozen all-zero initialization. + +``cached-partial-5hz`` + Keep completed one-second chunks in the official streaming cache. At + every non-boundary 200-ms cutoff, snapshot the complete MiniCPM streaming + state, append only the current causal partial AV tail with + ``is_last_chunk=True``, read its state, and restore the snapshot. At a + one-second boundary, append the complete chunk to the main cache. The + branch is admitted as ``mode=causal-prefix`` only after mandatory cache + replay, future-audio intervention, and native-boundary checks pass. + +The saved payload follows the existing global-feature contract consumed by +``easycom_head.py``: dense zero-based ``label_indices`` and ``states`` aligned +to the 5 Hz target rows. +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import math +from contextlib import AbstractContextManager +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch +from tqdm import tqdm + +import streaming_turn_model_directly as direct_media + + +MODEL_ID = "openbmb/MiniCPM-o-4_5" +TESTED_TRANSFORMERS_VERSION = "4.51.0" +FEATURE_FORMAT = "minicpmo-causal-global-v1" +EPSILON = 1e-7 + + +@dataclass(frozen=True) +class MiniCPMStreamConfig: + """Frozen extraction and verification settings.""" + + step_seconds: float = 0.2 + native_chunk_seconds: float = 1.0 + system_prompt: str = "Streaming Omni Conversation." + max_slice_nums: int = 1 + boundary_max_abs: float = 1e-3 + boundary_min_cosine: float = 0.9999 + replay_max_abs: float = 1e-3 + replay_min_cosine: float = 0.9999 + + def validate(self) -> None: + if self.step_seconds <= 0: + raise ValueError("step_seconds must be positive") + if self.native_chunk_seconds <= 0: + raise ValueError("native_chunk_seconds must be positive") + ratio = self.native_chunk_seconds / self.step_seconds + if not math.isclose(ratio, round(ratio), abs_tol=1e-9): + raise ValueError("native chunk / step must be integral") + if self.max_slice_nums <= 0: + raise ValueError("max_slice_nums must be positive") + for name, value in ( + ("boundary_max_abs", self.boundary_max_abs), + ("replay_max_abs", self.replay_max_abs), + ): + if value < 0: + raise ValueError(f"{name} must be non-negative") + for name, value in ( + ("boundary_min_cosine", self.boundary_min_cosine), + ("replay_min_cosine", self.replay_min_cosine), + ): + if not -1.0 <= value <= 1.0: + raise ValueError(f"{name} must be in [-1,1]") + + +@dataclass +class ExtractionResult: + times: torch.Tensor + states: torch.Tensor + state_source_times: torch.Tensor + native_update_mask: torch.Tensor + initialization_mask: torch.Tensor + verification: list[dict[str, Any]] + partial_tail_admissible: bool + + +def _time_grid(start: float, end: float, step: float) -> list[float]: + """Return right-edge cutoffs ``start + k*step`` through ``end``.""" + + count_f = (end - start) / step + count = int(round(count_f)) + if not math.isclose(count_f, count, rel_tol=0.0, abs_tol=1e-6): + raise ValueError( + f"duration must be an integer number of {step:g}s steps: " + f"start={start:g}, end={end:g}" + ) + return [start + (index + 1) * step for index in range(count)] + + +def _is_native_boundary( + current_time: float, + *, + origin: float, + chunk_seconds: float, +) -> bool: + units = (current_time - origin) / chunk_seconds + return math.isclose(units, round(units), rel_tol=0.0, abs_tol=1e-6) + + +def _comparison( + cached: torch.Tensor, + reference: torch.Tensor, + *, + kind: str, + current_time: float, + max_abs_tolerance: float, + min_cosine: float, +) -> dict[str, Any]: + if cached.shape != reference.shape: + raise RuntimeError( + f"{kind} shape mismatch at {current_time:g}s: " + f"{tuple(cached.shape)} vs {tuple(reference.shape)}" + ) + cached_f = cached.float().reshape(1, -1) + reference_f = reference.float().reshape(1, -1) + difference = (cached_f - reference_f).abs() + max_abs = float(difference.max().item()) if difference.numel() else 0.0 + mean_abs = float(difference.mean().item()) if difference.numel() else 0.0 + cached_norm = float(cached_f.norm().item()) + reference_norm = float(reference_f.norm().item()) + if torch.equal(cached_f, reference_f): + cosine = 1.0 + elif cached_norm == 0.0 and reference_norm == 0.0: + cosine = 1.0 + elif cached_norm == 0.0 or reference_norm == 0.0: + cosine = 0.0 + else: + cosine = float( + torch.nn.functional.cosine_similarity( + cached_f, reference_f, dim=-1 + ).item() + ) + passed = max_abs <= max_abs_tolerance and cosine >= min_cosine + return { + "kind": kind, + "time": float(current_time), + "max_abs": max_abs, + "mean_abs": mean_abs, + "cosine": cosine, + "max_abs_tolerance": float(max_abs_tolerance), + "min_cosine": float(min_cosine), + "passed": bool(passed), + } + + +class _FinalHiddenCapture(AbstractContextManager["_FinalHiddenCapture"]): + """Capture the final normalized Qwen3 state used by MiniCPM-o's LM head.""" + + def __init__(self, model: Any) -> None: + try: + self.module = model.llm.model.norm + except AttributeError as exc: + raise RuntimeError( + "MiniCPM-o readout contract changed: expected model.llm.model.norm" + ) from exc + self.value: torch.Tensor | None = None + self._handle: Any | None = None + + def __enter__(self) -> "_FinalHiddenCapture": + def hook(_module: Any, _inputs: Any, output: Any) -> None: + tensor = output[0] if isinstance(output, (tuple, list)) else output + if not isinstance(tensor, torch.Tensor) or tensor.ndim != 3: + raise RuntimeError( + "Unexpected MiniCPM-o final-norm output: " + f"{type(tensor)!r} shape={getattr(tensor, 'shape', None)}" + ) + self.value = tensor[:, -1, :].detach().float().cpu() + + self._handle = self.module.register_forward_hook(hook) + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + if self._handle is not None: + self._handle.remove() + self._handle = None + + +def validate_transformers_version( + version: str, + *, + allow_untested: bool = False, +) -> str: + if version != TESTED_TRANSFORMERS_VERSION and not allow_untested: + raise RuntimeError( + "MiniCPM-o 4.5 officially pins transformers==4.51.0; " + f"found {version}. Pass --allow-untested-transformers only for a " + "separate compatibility audit." + ) + return version + + +def _cache_length(cache: Any | None) -> int: + if cache is None: + return 0 + if hasattr(cache, "get_seq_length"): + return int(cache.get_seq_length()) + if hasattr(cache, "key_cache") and cache.key_cache: + return int(cache.key_cache[0].shape[-2]) + if isinstance(cache, (tuple, list)) and cache: + return int(cache[0][0].shape[-2]) + raise TypeError(f"Unsupported cache type: {type(cache)!r}") + + +def _cache_checksum(cache: Any | None) -> float | None: + """Match the first-key checksum stored by MiniCPM's snapshot object.""" + + if cache is None: + return None + if hasattr(cache, "self_attention_cache"): + cache = cache.self_attention_cache + if hasattr(cache, "key_cache") and cache.key_cache: + return float(cache.key_cache[0].sum().item()) + if isinstance(cache, (tuple, list)) and cache: + return float(cache[0][0].sum().item()) + raise TypeError(f"Unsupported cache type: {type(cache)!r}") + + +class MiniCPMInputOnlyEngine: + """Version-gated adapter around MiniCPM-o's simplex streaming prefill.""" + + def __init__( + self, + model: Any, + processor: Any, + config: MiniCPMStreamConfig, + *, + session_id: str, + ) -> None: + self.model = model + self.processor = processor + self.config = config + self.session_id = session_id + self.last_restore_audit: dict[str, Any] | None = None + try: + self.hidden_size = int(model.llm.config.hidden_size) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError("Cannot resolve MiniCPM-o LLM hidden size") from exc + for name in ( + "streaming_prefill", + "save_speculative_snapshot", + "restore_speculative_snapshot", + "reset_session", + ): + if not hasattr(model, name): + raise RuntimeError( + f"MiniCPM-o revision lacks required input-only API: {name}" + ) + + def reset(self) -> None: + self.model.reset_session(reset_token2wav_cache=False) + self.model.streaming_prefill( + session_id=self.session_id, + msgs=[ + { + "role": "system", + "content": [self.config.system_prompt], + } + ], + omni_mode=True, + max_slice_nums=self.config.max_slice_nums, + use_tts_template=False, + enable_thinking=False, + is_last_chunk=False, + processor=self.processor, + ) + + def _prefill( + self, + *, + frame: Any, + audio: np.ndarray, + is_last_chunk: bool, + ) -> torch.Tensor: + if not isinstance(audio, np.ndarray) or audio.ndim != 1 or not len(audio): + raise ValueError("MiniCPM-o audio must be a non-empty mono ndarray") + audio = np.asarray(audio, dtype=np.float32) + with _FinalHiddenCapture(self.model) as capture: + self.model.streaming_prefill( + session_id=self.session_id, + msgs=[{"role": "user", "content": [frame, audio]}], + omni_mode=True, + max_slice_nums=self.config.max_slice_nums, + use_tts_template=False, + enable_thinking=False, + is_last_chunk=is_last_chunk, + processor=self.processor, + ) + if capture.value is None: + raise RuntimeError("MiniCPM-o prefill produced no captured LLM state") + if capture.value.shape != (1, self.hidden_size): + raise RuntimeError( + "Unexpected MiniCPM-o state shape: " + f"{tuple(capture.value.shape)}, expected (1,{self.hidden_size})" + ) + if not torch.isfinite(capture.value).all(): + raise RuntimeError("MiniCPM-o state contains NaN/Inf") + return capture.value + + def commit_state(self, *, frame: Any, audio: np.ndarray) -> torch.Tensor: + return self._prefill(frame=frame, audio=audio, is_last_chunk=False) + + def temporary_state(self, *, frame: Any, audio: np.ndarray) -> torch.Tensor: + self.last_restore_audit = None + before_llm = _cache_length(getattr(self.model, "llm_past_key_values", None)) + before_audio = _cache_length( + getattr(self.model, "audio_past_key_values", None) + ) + before_chunk = int(getattr(self.model, "audio_chunk_idx", 0)) + before_history = tuple( + id(item) + for item in getattr(self.model, "_omni_chunk_history", ()) + ) + before_position = int( + getattr(self.model, "streaming_position_offset", 0) + ) + before_window_enabled = bool( + getattr(self.model, "streaming_window_enabled", False) + ) + snapshot = self.model.save_speculative_snapshot() + try: + # A speculative prefill must never invoke MiniCPM's destructive + # round-pruning path: the upstream snapshot truncates appended KV + # tokens but cannot recreate a round that pruning removed. Native + # committed chunks still retain the model's normal window policy. + if hasattr(self.model, "streaming_window_enabled"): + self.model.streaming_window_enabled = False + try: + state = self._prefill( + frame=frame, + audio=audio, + is_last_chunk=True, + ) + finally: + if hasattr(self.model, "streaming_window_enabled"): + self.model.streaming_window_enabled = before_window_enabled + finally: + restored = self.model.restore_speculative_snapshot(snapshot) + if not restored: + raise RuntimeError("MiniCPM-o failed to restore a partial-tail snapshot") + after_llm = _cache_length(getattr(self.model, "llm_past_key_values", None)) + after_audio = _cache_length( + getattr(self.model, "audio_past_key_values", None) + ) + after_chunk = int(getattr(self.model, "audio_chunk_idx", 0)) + after_history = tuple( + id(item) + for item in getattr(self.model, "_omni_chunk_history", ()) + ) + after_position = int( + getattr(self.model, "streaming_position_offset", 0) + ) + after_window_enabled = bool( + getattr(self.model, "streaming_window_enabled", False) + ) + after_llm_checksum = _cache_checksum( + getattr(self.model, "llm_past_key_values", None) + ) + after_audio_checksum = _cache_checksum( + getattr(self.model, "audio_past_key_values", None) + ) + expected_llm_checksum = getattr(snapshot, "llm_cache_checksum", None) + expected_audio_checksum = getattr(snapshot, "audio_cache_checksum", None) + structure_matches = ( + after_llm, + after_audio, + after_chunk, + after_history, + after_position, + after_window_enabled, + ) == ( + before_llm, + before_audio, + before_chunk, + before_history, + before_position, + before_window_enabled, + ) + checksum_matches = ( + after_llm_checksum == expected_llm_checksum + and after_audio_checksum == expected_audio_checksum + ) + self.last_restore_audit = { + "kind": "snapshot_restore", + "llm_cache_length": after_llm, + "audio_cache_length": after_audio, + "audio_chunk_index": after_chunk, + "streaming_position_offset": after_position, + "window_pruning_suspended": bool(before_window_enabled), + "window_policy_restored": bool( + after_window_enabled == before_window_enabled + ), + "llm_checksum_matches": bool( + after_llm_checksum == expected_llm_checksum + ), + "audio_checksum_matches": bool( + after_audio_checksum == expected_audio_checksum + ), + "history_matches": bool(after_history == before_history), + "passed": bool(structure_matches and checksum_matches), + } + if not self.last_restore_audit["passed"]: + raise RuntimeError( + "Partial-tail restore changed committed MiniCPM-o state: " + f"{self.last_restore_audit}" + ) + return state + + +class MiniCPMFeatureExtractor: + """Schedule native commits and temporary tails on a causal AV reader.""" + + def __init__( + self, + reader: Any, + engine: Any, + config: MiniCPMStreamConfig, + ) -> None: + config.validate() + self.reader = reader + self.engine = engine + self.config = config + + def _frame_audio(self, start: float, end: float) -> tuple[Any, np.ndarray]: + # Audio uses [start,end). Excluding a source frame exactly at ``end`` + # gives the same half-open cutoff contract as the Qwen extractor. + # CausalSourceAV admits an EPSILON tolerance, so nextafter alone is not + # enough to exclude a frame whose PTS is exactly the right edge. + frame_time = max(start, end - max(4 * direct_media.EPSILON, 1e-6)) + frame = self.reader.frame_at_or_before(frame_time) + audio = self.reader.audio_slice(start, end) + return frame, audio + + def _commit(self, start: float, end: float) -> torch.Tensor: + frame, audio = self._frame_audio(start, end) + return self.engine.commit_state(frame=frame, audio=audio) + + def _temporary(self, start: float, end: float) -> torch.Tensor: + frame, audio = self._frame_audio(start, end) + return self.engine.temporary_state(frame=frame, audio=audio) + + def _replay_to( + self, + *, + origin: float, + current_time: float, + ) -> tuple[torch.Tensor, float]: + """Cleanly replay completed chunks and the optional current tail.""" + + self.engine.reset() + cursor = origin + state: torch.Tensor | None = None + while cursor + self.config.native_chunk_seconds <= current_time + EPSILON: + next_cursor = cursor + self.config.native_chunk_seconds + state = self._commit(cursor, next_cursor) + cursor = next_cursor + if cursor < current_time - EPSILON: + state = self._temporary(cursor, current_time) + if state is None: + state = torch.zeros(1, self.engine.hidden_size, dtype=torch.float32) + return state, cursor + + def extract_native_hold( + self, + *, + start: float, + end: float, + show_progress: bool, + ) -> ExtractionResult: + times = _time_grid(start, end, self.config.step_seconds) + self.engine.reset() + latest = torch.zeros(1, self.engine.hidden_size, dtype=torch.float32) + source_time = start + chunk_start = start + states: list[torch.Tensor] = [] + source_times: list[float] = [] + updates: list[bool] = [] + initialization: list[bool] = [] + with tqdm( + total=len(times), + desc="MiniCPM native 1Hz", + unit="state", + disable=not show_progress, + dynamic_ncols=True, + ) as progress: + for current_time in times: + update = _is_native_boundary( + current_time, + origin=start, + chunk_seconds=self.config.native_chunk_seconds, + ) + if update: + latest = self._commit(chunk_start, current_time) + chunk_start = current_time + source_time = current_time + states.append(latest.cpu()) + source_times.append(source_time) + updates.append(update) + initialization.append(source_time <= start + EPSILON) + progress.update(1) + return ExtractionResult( + times=torch.tensor(times, dtype=torch.float64), + states=torch.cat(states, dim=0), + state_source_times=torch.tensor(source_times, dtype=torch.float64), + native_update_mask=torch.tensor(updates, dtype=torch.bool), + initialization_mask=torch.tensor(initialization, dtype=torch.bool), + verification=[], + partial_tail_admissible=True, + ) + + def extract_cached_partial( + self, + *, + start: float, + end: float, + verify_every: int, + future_intervention_every: int, + allow_unverified: bool, + show_progress: bool, + ) -> ExtractionResult: + if verify_every < 0 or future_intervention_every < 0: + raise ValueError("verification intervals must be non-negative") + times = _time_grid(start, end, self.config.step_seconds) + self.engine.reset() + chunk_start = start + states: list[torch.Tensor] = [] + updates: list[bool] = [] + checks: list[dict[str, Any]] = [] + + with tqdm( + total=len(times), + desc="MiniCPM cached partial 5Hz", + unit="state", + disable=not show_progress, + dynamic_ncols=True, + ) as progress: + for step_index, current_time in enumerate(times, start=1): + boundary = _is_native_boundary( + current_time, + origin=start, + chunk_seconds=self.config.native_chunk_seconds, + ) + if boundary: + # Compare a temporary closed view of exactly the same full + # chunk with the native non-final commit before admitting + # this 5 Hz representation. + temporary = self._temporary(chunk_start, current_time) + restore_audit = getattr( + self.engine, "last_restore_audit", None + ) + if restore_audit is not None: + checks.append( + {**restore_audit, "time": float(current_time)} + ) + state = self._commit(chunk_start, current_time) + comparison = _comparison( + state, + temporary, + kind="native_boundary_equivalence", + current_time=current_time, + max_abs_tolerance=self.config.boundary_max_abs, + min_cosine=self.config.boundary_min_cosine, + ) + checks.append(comparison) + if not comparison["passed"] and not allow_unverified: + raise RuntimeError( + "MiniCPM partial-tail/native boundary check failed " + f"at {current_time:g}s: {comparison}" + ) + chunk_start = current_time + else: + state = self._temporary(chunk_start, current_time) + restore_audit = getattr( + self.engine, "last_restore_audit", None + ) + if restore_audit is not None: + checks.append( + {**restore_audit, "time": float(current_time)} + ) + + # Always audit the first row. Positive intervals add sampled + # checks without turning a full session into O(T^2) replay. + replay_due = step_index == 1 or ( + verify_every > 0 and step_index % verify_every == 0 + ) + if replay_due: + reference, rebuilt_until = self._replay_to( + origin=start, + current_time=current_time, + ) + comparison = _comparison( + state, + reference, + kind="clean_stream_replay", + current_time=current_time, + max_abs_tolerance=self.config.replay_max_abs, + min_cosine=self.config.replay_min_cosine, + ) + comparison["rebuilt_committed_until"] = float(rebuilt_until) + checks.append(comparison) + if not comparison["passed"] and not allow_unverified: + raise RuntimeError( + "MiniCPM clean replay check failed at " + f"{current_time:g}s: {comparison}" + ) + + future_due = ( + current_time < end - EPSILON + and ( + step_index == 1 + or ( + future_intervention_every > 0 + and step_index % future_intervention_every == 0 + ) + ) + ) + if future_due: + first = int( + math.floor( + current_time * direct_media.SAMPLE_RATE + EPSILON + ) + ) + last = min( + len(self.reader.audio), + first + direct_media.SAMPLE_RATE, + ) + original = self.reader.audio[first:last].copy() + if last > first: + self.reader.audio[first:last] = np.linspace( + -0.25, + 0.25, + last - first, + dtype=np.float32, + ) + try: + intervened, _ = self._replay_to( + origin=start, + current_time=current_time, + ) + finally: + self.reader.audio[first:last] = original + comparison = _comparison( + state, + intervened, + kind="future_audio_intervention", + current_time=current_time, + max_abs_tolerance=0.0, + min_cosine=1.0, + ) + checks.append(comparison) + if not comparison["passed"] and not allow_unverified: + raise RuntimeError( + "Future audio changed a causal MiniCPM state at " + f"{current_time:g}s: {comparison}" + ) + + states.append(state.cpu()) + updates.append(boundary) + progress.update(1) + + kinds = {check["kind"] for check in checks if check.get("passed")} + mandatory = { + "native_boundary_equivalence", + "clean_stream_replay", + "future_audio_intervention", + "snapshot_restore", + } + admissible = mandatory.issubset(kinds) and all( + bool(check.get("passed")) for check in checks + ) + if not admissible and not allow_unverified: + missing = sorted(mandatory - kinds) + raise RuntimeError( + "MiniCPM partial-tail branch lacks mandatory verification; " + f"missing={missing}. Run a >=1.2s smoke test or pass " + "--allow-unverified-partial-tail to save a deliberately " + "non-loadable audit artifact." + ) + return ExtractionResult( + times=torch.tensor(times, dtype=torch.float64), + states=torch.cat(states, dim=0), + state_source_times=torch.tensor(times, dtype=torch.float64), + native_update_mask=torch.tensor(updates, dtype=torch.bool), + initialization_mask=torch.zeros(len(times), dtype=torch.bool), + verification=checks, + partial_tail_admissible=admissible, + ) + + +def load_minicpmo( + model_id: str, + *, + revision: str | None, + dtype: str, + device: str, + attention: str, + allow_untested_transformers: bool, +) -> tuple[Any, Any, dict[str, Any]]: + """Load the frozen AV backbone without the TTS module.""" + + installed = importlib.metadata.version("transformers") + validate_transformers_version( + installed, + allow_untested=allow_untested_transformers, + ) + from transformers import AutoModel, AutoProcessor + + dtype_value: Any = { + "auto": "auto", + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + }[dtype] + common: dict[str, Any] = { + "trust_remote_code": True, + } + if revision is not None: + common["revision"] = revision + processor = AutoProcessor.from_pretrained(model_id, **common) + model = AutoModel.from_pretrained( + model_id, + torch_dtype=dtype_value, + attn_implementation=attention, + init_vision=True, + init_audio=True, + init_tts=False, + device_map={"": device}, + **common, + ) + model.eval() + model.requires_grad_(False) + resolved_revision = getattr(model.config, "_commit_hash", None) + metadata = { + "transformers_version": installed, + "requested_revision": revision, + "resolved_revision": resolved_revision, + "attention_implementation": attention, + "dtype": dtype, + "device": device, + "init_tts": False, + } + return model, processor, metadata + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--video", type=Path) + source.add_argument("--easycom-session-dir", type=Path) + source.add_argument("--easycom-manifest", type=Path) + parser.add_argument("--session") + parser.add_argument("--audio", type=Path) + parser.add_argument("--crop") + parser.add_argument("--target", type=Path) + parser.add_argument("--output", type=Path) + parser.add_argument( + "--mode", + choices=("native-1hz-hold", "cached-partial-5hz"), + required=True, + ) + parser.add_argument("--model-id", default=MODEL_ID) + parser.add_argument( + "--revision", + help="Pinned Hugging Face revision; strongly recommended for paper runs", + ) + parser.add_argument("--start", type=float, default=0.0) + parser.add_argument("--end", type=float) + parser.add_argument("--step-seconds", type=float, default=0.2) + parser.add_argument("--native-chunk-seconds", type=float, default=1.0) + parser.add_argument( + "--system-prompt", + default="Streaming Omni Conversation.", + ) + parser.add_argument("--max-slice-nums", type=int, default=1) + parser.add_argument("--boundary-max-abs", type=float, default=1e-3) + parser.add_argument("--boundary-min-cosine", type=float, default=0.9999) + parser.add_argument("--replay-max-abs", type=float, default=1e-3) + parser.add_argument("--replay-min-cosine", type=float, default=0.9999) + parser.add_argument( + "--verify-every", + type=int, + default=0, + help="Additional clean replays every N rows; the first row is always checked", + ) + parser.add_argument( + "--future-intervention-every", + type=int, + default=0, + help=( + "Additional future-audio checks every N rows; the first eligible " + "row is always checked" + ), + ) + parser.add_argument("--allow-unverified-partial-tail", action="store_true") + parser.add_argument("--frame-cache-size", type=int, default=8) + parser.add_argument( + "--decoder", + choices=("decord", "torchcodec", "pynvcodec"), + default="decord", + ) + parser.add_argument("--decoder-gpu-id", type=int, default=0) + parser.add_argument("--decoder-threads", type=int, default=1) + parser.add_argument("--max-open-readers", type=int, default=2) + parser.add_argument( + "--dtype", + choices=("auto", "bfloat16", "float16", "float32"), + default="bfloat16", + ) + parser.add_argument( + "--attention", + choices=("sdpa", "flash_attention_2"), + default="sdpa", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--allow-untested-transformers", action="store_true") + parser.add_argument("--media-only", action="store_true") + parser.add_argument("--no-progress", action="store_true") + return parser.parse_args(argv) + + +def _validate_args(args: argparse.Namespace) -> None: + if args.output is None and not args.media_only: + raise SystemExit("--output is required unless --media-only is used") + if args.audio is not None and args.video is None: + raise SystemExit("--audio is supported only with --video") + if args.session is not None and args.easycom_manifest is None: + raise SystemExit("--session is used only with --easycom-manifest") + if args.easycom_manifest is not None and args.session is None: + raise SystemExit("--easycom-manifest requires --session Session_N") + if args.frame_cache_size < 0: + raise SystemExit("--frame-cache-size must be non-negative") + if args.verify_every < 0 or args.future_intervention_every < 0: + raise SystemExit("verification intervals must be non-negative") + if not args.media_only and args.revision is None: + print( + "WARNING: --revision is not pinned; this run is not paper-reproducible.", + flush=True, + ) + + +def _build_reader( + args: argparse.Namespace, +) -> tuple[Any, Path | None, int | None, float, str, Any]: + manifest_record = None + manifest_target = None + if args.video is not None: + specs = [ + direct_media.SourceClipSpec( + path=args.video, + crop=direct_media._parse_crop(args.crop), + audio_path=args.audio, + ) + ] + source_kind = "single_original_video" + elif args.easycom_session_dir is not None: + if args.crop is not None: + raise SystemExit( + "EasyCom direct extraction uses the fixed rightmost-1920 crop" + ) + specs = direct_media._easycom_directory_specs(args.easycom_session_dir) + source_kind = "easycom_dense_original_clips" + else: + if args.crop is not None: + raise SystemExit("EasyCom manifest supplies crop metadata") + specs, manifest_record, manifest_target = direct_media._manifest_session_specs( + args.easycom_manifest.expanduser().resolve(), args.session + ) + source_kind = "easycom_manifest_original_clips" + + target_path, target_steps, end = direct_media._resolve_duration_and_target( + args, manifest_target + ) + reader = direct_media.CausalSourceAV( + specs, + output_fps=1.0 / args.step_seconds, + required_duration=end, + frame_cache_size=args.frame_cache_size, + decoder_backend=args.decoder, + decoder_gpu_id=args.decoder_gpu_id, + decoder_threads=args.decoder_threads, + max_open_readers=args.max_open_readers, + ) + is_easycom = ( + args.easycom_session_dir is not None + or args.easycom_manifest is not None + ) + if is_easycom: + for clip in reader.clips: + if clip.height != direct_media.EASYCOM_HEIGHT: + raise RuntimeError( + f"Expected EasyCom height {direct_media.EASYCOM_HEIGHT}, " + f"got {clip.height}: {clip.path}" + ) + if not math.isclose(clip.source_fps, 20.0, abs_tol=1e-6): + raise RuntimeError( + f"Expected EasyCom source FPS 20, got {clip.source_fps:g}: " + f"{clip.path}" + ) + if target_steps is not None: + source_frames = sum(clip.num_frames for clip in reader.clips) + media_steps = source_frames // 4 + if media_steps != target_steps: + raise RuntimeError( + f"EasyCom source gives {media_steps} 5-Hz rows, target has " + f"{target_steps}" + ) + direct_media._print_media_summary(reader, args.start, end) + return reader, target_path, target_steps, end, source_kind, manifest_record + + +def main( + argv: Sequence[str] | None = None, + *, + model: Any | None = None, + processor: Any | None = None, +) -> None: + args = parse_args(argv) + _validate_args(args) + if (model is None) != (processor is None): + raise ValueError("model and processor must be supplied together") + config = MiniCPMStreamConfig( + step_seconds=args.step_seconds, + native_chunk_seconds=args.native_chunk_seconds, + system_prompt=args.system_prompt, + max_slice_nums=args.max_slice_nums, + boundary_max_abs=args.boundary_max_abs, + boundary_min_cosine=args.boundary_min_cosine, + replay_max_abs=args.replay_max_abs, + replay_min_cosine=args.replay_min_cosine, + ) + config.validate() + if not math.isclose(args.native_chunk_seconds, 1.0, abs_tol=1e-9): + raise SystemExit( + "MiniCPM-o 4.5 native streaming is fixed at 1.0s; " + "--native-chunk-seconds is exposed only to reject protocol drift" + ) + + ( + reader, + target_path, + target_steps, + end, + source_kind, + manifest_record, + ) = _build_reader(args) + if args.media_only: + probe_end = min(end, args.start + args.step_seconds) + reader.frame_at_or_before(float(np.nextafter(probe_end, -np.inf))) + reader.audio_slice(args.start, probe_end) + print(f"Decoded MiniCPM media probe successfully with {args.decoder}") + return + + model_metadata: dict[str, Any] + if model is None: + model, processor, model_metadata = load_minicpmo( + args.model_id, + revision=args.revision, + dtype=args.dtype, + device=args.device, + attention=args.attention, + allow_untested_transformers=args.allow_untested_transformers, + ) + else: + model.eval() + model.requires_grad_(False) + model_metadata = { + "transformers_version": importlib.metadata.version("transformers"), + "requested_revision": args.revision, + "resolved_revision": getattr(model.config, "_commit_hash", None), + "preloaded": True, + } + + session_name = args.session or ( + args.easycom_session_dir.name + if args.easycom_session_dir is not None + else args.video.stem + ) + engine = MiniCPMInputOnlyEngine( + model, + processor, + config, + session_id=f"feature-extract-{session_name}", + ) + extractor = MiniCPMFeatureExtractor(reader, engine, config) + if args.mode == "native-1hz-hold": + result = extractor.extract_native_hold( + start=args.start, + end=end, + show_progress=not args.no_progress, + ) + else: + result = extractor.extract_cached_partial( + start=args.start, + end=end, + verify_every=args.verify_every, + future_intervention_every=args.future_intervention_every, + allow_unverified=args.allow_unverified_partial_tail, + show_progress=not args.no_progress, + ) + + if target_steps is not None and len(result.times) != target_steps: + raise RuntimeError( + f"Extractor returned {len(result.times)} states, expected {target_steps}" + ) + loadable_mode = ( + "causal-prefix" + if result.partial_tail_admissible + else "minicpmo-unverified-partial-tail" + ) + payload = { + "format": FEATURE_FORMAT, + "model_id": args.model_id, + "mode": loadable_mode, + "extractor_mode": args.mode, + "stream_representation": ( + "native_1s_input_only_state_forward_held_to_5hz" + if args.mode == "native-1hz-hold" + else "committed_1s_cache_plus_temporary_causal_partial_tail_at_5hz" + ), + "config": asdict(config), + "model_metadata": model_metadata, + "strict_causal": True, + "generation_called": False, + "tts_loaded": False, + "partial_tail_admissible": result.partial_tail_admissible, + "initialization_policy": ( + "zeros_until_first_completed_native_chunk" + if args.mode == "native-1hz-hold" + else None + ), + "start": args.start, + "end": end, + "num_steps": len(result.times), + "label_indices": torch.arange(len(result.times), dtype=torch.long), + "valid_10_horizon_steps": max(0, len(result.times) - 10), + "times": result.times, + "state_source_times": result.state_source_times, + "native_update_mask": result.native_update_mask, + "initialization_mask": result.initialization_mask, + "states": result.states, + "verification": result.verification, + "target_path": str(target_path) if target_path is not None else None, + "source_kind": source_kind, + "source_manifest_record": manifest_record, + "direct_media": reader.metadata(), + "video_sampling_grid": "latest_source_pts_strictly_before_each_cutoff", + "audio_interval": "half_open_[start,end)", + } + args.output.parent.mkdir(parents=True, exist_ok=True) + torch.save(payload, args.output) + print(f"Saved {len(result.times)} MiniCPM states to {args.output}") + print(f"states: {tuple(result.states.shape)}") + print(f"mode: {loadable_mode}; extractor_mode: {args.mode}") + if result.verification: + kinds = sorted({item["kind"] for item in result.verification}) + for kind in kinds: + items = [item for item in result.verification if item["kind"] == kind] + worst_abs = max(float(item["max_abs"]) for item in items) + worst_cosine = min(float(item["cosine"]) for item in items) + print( + f" {kind}: n={len(items)}, max_abs={worst_abs:.6g}, " + f"min_cosine={worst_cosine:.9f}" + ) + + +if __name__ == "__main__": + main() diff --git a/test_streaming_minicpmo_features_directly.py b/test_streaming_minicpmo_features_directly.py new file mode 100644 index 0000000..9b806f7 --- /dev/null +++ b/test_streaming_minicpmo_features_directly.py @@ -0,0 +1,230 @@ +"""CPU-only contract tests for MiniCPM-o streaming feature schedules.""" + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import torch + +from easycom_head import load_session_data +from local_turn_features import reference_alignment +from streaming_minicpmo_features_directly import ( + FEATURE_FORMAT, + MiniCPMFeatureExtractor, + MiniCPMStreamConfig, + _comparison, + _time_grid, + validate_transformers_version, +) + + +class FakeReader: + def __init__(self, duration: float = 3.0) -> None: + self.audio = np.arange( + int(duration * 16_000), dtype=np.float32 + ) / 16_000.0 + + def frame_at_or_before(self, timestamp: float) -> float: + return float(timestamp) + + def audio_slice(self, start: float, end: float) -> np.ndarray: + first = int(round(start * 16_000)) + last = int(round(end * 16_000)) + return self.audio[first:last].copy() + + +class FakeEngine: + hidden_size = 3 + + def __init__(self) -> None: + self.committed_until = 0.0 + self.reset_count = 0 + self.last_restore_audit = None + + def reset(self) -> None: + self.committed_until = 0.0 + self.reset_count += 1 + + @staticmethod + def _state(frame: float, audio: np.ndarray) -> torch.Tensor: + return torch.tensor( + [[float(audio[-1]), float(frame), float(len(audio))]], + dtype=torch.float32, + ) + + def commit_state(self, *, frame: float, audio: np.ndarray) -> torch.Tensor: + self.committed_until += len(audio) / 16_000.0 + return self._state(frame, audio) + + def temporary_state(self, *, frame: float, audio: np.ndarray) -> torch.Tensor: + before = self.committed_until + state = self._state(frame, audio) + if self.committed_until != before: + raise AssertionError("temporary state modified the fake main cache") + self.last_restore_audit = { + "kind": "snapshot_restore", + "llm_cache_length": 0, + "audio_cache_length": 0, + "audio_chunk_index": 0, + "streaming_position_offset": 0, + "window_pruning_suspended": True, + "window_policy_restored": True, + "llm_checksum_matches": True, + "audio_checksum_matches": True, + "history_matches": True, + "passed": True, + } + return state + + +class MiniCPMStreamingTests(unittest.TestCase): + def setUp(self) -> None: + self.config = MiniCPMStreamConfig( + boundary_max_abs=0.0, + boundary_min_cosine=1.0, + replay_max_abs=0.0, + replay_min_cosine=1.0, + ) + + def test_time_grid_is_right_edged_and_rejects_fractional_tail(self) -> None: + self.assertEqual(_time_grid(0.0, 1.0, 0.2), [0.2, 0.4, 0.6, 0.8, 1.0]) + with self.assertRaisesRegex(ValueError, "integer number"): + _time_grid(0.0, 1.1, 0.2) + + def test_native_hold_has_no_future_backfill(self) -> None: + extractor = MiniCPMFeatureExtractor(FakeReader(), FakeEngine(), self.config) + result = extractor.extract_native_hold( + start=0.0, + end=2.0, + show_progress=False, + ) + self.assertEqual(tuple(result.states.shape), (10, 3)) + self.assertTrue(torch.equal(result.states[:4], torch.zeros(4, 3))) + self.assertTrue(result.initialization_mask[:4].all()) + self.assertFalse(result.initialization_mask[4:].any()) + self.assertEqual( + result.native_update_mask.nonzero().flatten().tolist(), + [4, 9], + ) + self.assertTrue(torch.equal(result.states[4], result.states[8])) + self.assertFalse(torch.equal(result.states[8], result.states[9])) + self.assertEqual( + result.state_source_times.tolist(), + [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0], + ) + + def test_cached_partial_updates_every_row_and_passes_audits(self) -> None: + reader = FakeReader() + extractor = MiniCPMFeatureExtractor(reader, FakeEngine(), self.config) + original_audio = reader.audio.copy() + result = extractor.extract_cached_partial( + start=0.0, + end=2.0, + verify_every=0, + future_intervention_every=0, + allow_unverified=False, + show_progress=False, + ) + self.assertTrue(result.partial_tail_admissible) + self.assertTrue(np.array_equal(reader.audio, original_audio)) + self.assertFalse(result.initialization_mask.any()) + self.assertTrue(torch.equal(result.state_source_times, result.times)) + self.assertEqual( + result.native_update_mask.nonzero().flatten().tolist(), + [4, 9], + ) + self.assertEqual( + {item["kind"] for item in result.verification}, + { + "native_boundary_equivalence", + "clean_stream_replay", + "future_audio_intervention", + "snapshot_restore", + }, + ) + self.assertTrue(all(item["passed"] for item in result.verification)) + self.assertFalse(torch.equal(result.states[0], result.states[1])) + self.assertLess(float(result.states[0, 1]), 0.2) + + def test_short_partial_run_is_rejected_without_boundary_audit(self) -> None: + extractor = MiniCPMFeatureExtractor(FakeReader(), FakeEngine(), self.config) + with self.assertRaisesRegex(RuntimeError, "missing=.*native_boundary"): + extractor.extract_cached_partial( + start=0.0, + end=0.8, + verify_every=0, + future_intervention_every=0, + allow_unverified=False, + show_progress=False, + ) + + def test_comparison_rejects_different_states(self) -> None: + item = _comparison( + torch.zeros(1, 2), + torch.ones(1, 2), + kind="test", + current_time=0.2, + max_abs_tolerance=0.0, + min_cosine=1.0, + ) + self.assertFalse(item["passed"]) + self.assertEqual(item["max_abs"], 1.0) + + def test_transformers_version_is_explicitly_gated(self) -> None: + self.assertEqual(validate_transformers_version("4.51.0"), "4.51.0") + with self.assertRaisesRegex(RuntimeError, "found 4.52.0"): + validate_transformers_version("4.52.0") + self.assertEqual( + validate_transformers_version("4.52.0", allow_untested=True), + "4.52.0", + ) + + def test_saved_payload_loads_in_existing_head_and_unverified_fails(self) -> None: + states = torch.randn(20, 8) + payload = { + "format": FEATURE_FORMAT, + "states": states, + "times": (torch.arange(20, dtype=torch.float64) + 1) / 5, + "label_indices": torch.arange(20), + "mode": "causal-prefix", + "extractor_mode": "native-1hz-hold", + "config": {"step_seconds": 0.2}, + } + targets = np.tile( + np.asarray( + [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], + dtype=np.float32, + ), + (5, 1), + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + feature_root = root / "features" + target_root = root / "targets" + feature_root.mkdir() + target_root.mkdir() + feature_path = feature_root / "Session_4.pt" + torch.save(payload, feature_path) + np.save(target_root / "Session_4.npy", targets) + loaded = load_session_data(feature_root, target_root, 4) + self.assertTrue(torch.equal(loaded.states, states)) + times, indices = reference_alignment(payload) + self.assertTrue(torch.equal(times, payload["times"])) + self.assertTrue(torch.equal(indices, payload["label_indices"])) + + bad_reference = dict(payload) + bad_reference["times"] = payload["times"] + 0.2 + with self.assertRaisesRegex(ValueError, "Reference times"): + reference_alignment(bad_reference) + + payload["mode"] = "minicpmo-unverified-partial-tail" + torch.save(payload, feature_path) + with self.assertRaisesRegex(ValueError, "expected 'causal-prefix'"): + load_session_data(feature_root, target_root, 4) + + +if __name__ == "__main__": + unittest.main()