icassp2027-experiment/test_streaming_minicpmo_features_directly.py

"""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()