"""CPU-only schedule and feature-contract tests for Qwen chunk ablations.""" from __future__ import annotations import tempfile import types import unittest from pathlib import Path from unittest import mock import numpy as np import torch import streaming_turn_model as core from easycom_head import load_session_data from local_turn_features import reference_alignment class FakeModel: def __init__(self) -> None: self.config = types.SimpleNamespace( seconds_per_chunk=2.0, position_id_per_seconds=25.0, text_config=types.SimpleNamespace(hidden_size=3), ) self.dtype = torch.float32 self.device = torch.device("cpu") class FakeCache: def __init__(self, count: int) -> None: self.count = count def fake_pieces() -> core.OpenMediaPieces: return core.OpenMediaPieces( prefix_embeddings=torch.zeros(1, 1, 3), prefix_position_ids=torch.zeros(3, 1, 1, dtype=torch.long), payload_embeddings=torch.zeros(1, 1, 3), payload_position_ids=torch.zeros(3, 1, 1, dtype=torch.long), closure_ids=torch.zeros(1, 1, dtype=torch.long), closure_position_ids=torch.ones(3, 1, 1, dtype=torch.long), ) class RecordingMixin: def _prepare_pieces(self, **kwargs): # type: ignore[no-untyped-def] self.prepared_intervals.append(tuple(kwargs["interval"])) return fake_pieces() class RecordingHold(RecordingMixin, core.CompletedBlockHoldExtractor): def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def] super().__init__(*args, **kwargs) self.prepared_intervals: list[tuple[float, float]] = [] class RecordingMicro(RecordingMixin, core.MicroChunkReplayExtractor): def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def] super().__init__(*args, **kwargs) self.prepared_intervals: list[tuple[float, float]] = [] def fake_append( # type: ignore[no-untyped-def] model, pieces, *, past_key_values, rope_offset, include_prefix, ): count = 0 if past_key_values is None else past_key_values.count return FakeCache(count + 1), rope_offset + 1, pieces.closure_position_ids def fake_read(model, cache, **kwargs): # type: ignore[no-untyped-def] return torch.full((1, 3), float(cache.count), dtype=torch.float32) class QwenChunkAblationTests(unittest.TestCase): def setUp(self) -> None: self.model = FakeModel() self.processor = object() self.config = core.StreamConfig(memory_seconds=4.0) self.patches = [ mock.patch.object( core, "_readout_query_ids", return_value=torch.zeros(1, 1, dtype=torch.long), ), mock.patch.object(core, "_append_open_pieces", side_effect=fake_append), mock.patch.object(core, "_read_open_state", side_effect=fake_read), ] for patcher in self.patches: patcher.start() self.addCleanup(patcher.stop) def test_completed_native_blocks_are_held_without_future_backfill(self) -> None: extractor = RecordingHold(self.model, self.processor, self.config) result = extractor.extract( video_path="unused.mp4", audio_path=None, start=0.0, end=4.0, use_audio_in_video=True, show_progress=False, ) self.assertEqual(tuple(result.states.shape), (20, 3)) self.assertTrue(torch.equal(result.states[:9], torch.zeros(9, 3))) self.assertTrue(torch.equal(result.states[9:19], torch.ones(10, 3))) self.assertTrue(torch.equal(result.states[19], torch.full((3,), 2.0))) self.assertEqual( result.diagnostics["state_update_mask"].nonzero().flatten().tolist(), [9, 19], ) self.assertTrue( torch.equal( result.diagnostics["state_update_mask"], result.diagnostics["payload_commit_mask"], ) ) self.assertTrue( torch.equal( result.diagnostics["state_update_mask"], result.diagnostics["native_boundary_mask"], ) ) self.assertTrue(result.diagnostics["initialization_mask"][:9].all()) self.assertFalse(result.diagnostics["initialization_mask"][9:].any()) self.assertEqual( result.diagnostics["state_source_times"].tolist(), [0.0] * 9 + [2.0] * 10 + [4.0], ) self.assertEqual(extractor.prepared_intervals, [(0.0, 2.0), (2.0, 4.0)]) self.assertEqual(result.diagnostics["production_processor_calls"], 2) self.assertAlmostEqual( result.diagnostics["production_prepared_media_seconds"], 4.0 ) def test_micro_chunks_update_every_200ms_without_growing_tail(self) -> None: config = core.StreamConfig(memory_seconds=2.0) extractor = RecordingMicro(self.model, self.processor, config) result = extractor.extract( video_path="unused.mp4", audio_path=None, start=0.0, end=2.0, use_audio_in_video=True, show_progress=False, ) expected = torch.arange(1, 11, dtype=torch.float32).view(-1, 1) expected = expected.expand(-1, 3) self.assertTrue(torch.equal(result.states, expected)) self.assertTrue(result.diagnostics["state_update_mask"].all()) self.assertTrue(result.diagnostics["payload_commit_mask"].all()) self.assertEqual( result.diagnostics["native_boundary_mask"] .nonzero() .flatten() .tolist(), [9], ) self.assertTrue( torch.equal( result.diagnostics["state_source_times"], result.times ) ) self.assertEqual(len(extractor.prepared_intervals), 10) for index, interval in enumerate(extractor.prepared_intervals): self.assertAlmostEqual(interval[0], index * 0.2) self.assertAlmostEqual(interval[1], (index + 1) * 0.2) self.assertEqual(result.diagnostics["production_processor_calls"], 10) self.assertAlmostEqual( result.diagnostics["production_prepared_media_seconds"], 2.0 ) self.assertEqual(result.diagnostics["tmrope_positions_per_micro_step"], 5) def test_micro_tmrope_step_must_tile_the_native_chunk(self) -> None: self.assertEqual(core._tmrope_positions_per_step(self.model, 0.2), 5) with self.assertRaisesRegex(ValueError, "Non-integral TMRoPE"): core._tmrope_positions_per_step(self.model, 0.15) def test_both_modes_load_through_head_and_local_alignment_contracts(self) -> None: states = torch.randn(20, 8) 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" np.save(target_root / "Session_4.npy", targets) for mode in ("completed-block-hold", "micro-chunk-200ms"): payload = { "states": states, "times": (torch.arange(20, dtype=torch.float64) + 1) / 5, "label_indices": torch.arange(20), "mode": mode, "config": {"step_seconds": 0.2}, "start": 0.0, } torch.save(payload, feature_path) loaded = load_session_data( feature_root, target_root, 4, expected_mode=mode, ) 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"])) with self.assertRaisesRegex(ValueError, "expected 'causal-prefix'"): load_session_data(feature_root, target_root, 4) if __name__ == "__main__": unittest.main()