"""CPU cache semantics, output routing and both feature-loader contracts.""" from __future__ import annotations import copy import json import tempfile import types import unittest from pathlib import Path from unittest import mock import numpy as np import torch import extract_easycom_multiwindow as driver import streaming_turn_model as core import streaming_turn_model_directly as direct from easycom_head import load_session_data from local_turn_features import LOCAL_FORMAT, load_local_features, reference_alignment from streaming_turn_multiwindow import DEFAULT_MEMORY_SECONDS, MultiWindowOpenStreamReplayExtractor, validate_memory_seconds class FakeModel: config = types.SimpleNamespace(seconds_per_chunk=2.0, position_id_per_seconds=25.0, text_config=types.SimpleNamespace(hidden_size=3)) device = torch.device("cpu") dtype = torch.float32 class FakeCache: def __init__(self): self.tokens = [] def append(model, pieces, *, past_key_values, rope_offset, include_prefix): cache = FakeCache() if past_key_values is None else past_key_values if include_prefix != (not cache.tokens): raise AssertionError("prefix ownership mismatch or readout committed") values = pieces.payload_embeddings.flatten().tolist() # Mutation exercises clone_cache isolation, not a conveniently immutable fake. cache.tokens.append((values, rope_offset)) return cache, rope_offset + 1, pieces.closure_position_ids def read(model, cache, *, last_n_layers=1, layer_fusion="mean", **kwargs): value = 0.0 for values, offset in cache.tokens: value = value * 1.17 + sum(values) + offset / 100 state = torch.tensor([[value, cache.tokens[0][0][0], cache.tokens[-1][0][1]]]) if layer_fusion == "stack": return torch.stack([state + layer for layer in range(last_n_layers)], dim=1) return state class MultiwindowTests(unittest.TestCase): def setUp(self): self.prepared = [] self.future_change = False def prepare(extractor, *, interval, **kwargs): start, end = interval self.prepared.append(interval) signal = 1000.0 if self.future_change and end > 13.0 else 0.0 return core.OpenMediaPieces( prefix_embeddings=torch.zeros(1, 1, 3), prefix_position_ids=torch.zeros(3, 1, 1, dtype=torch.long), payload_embeddings=torch.tensor([[[start, end, signal]]]), payload_position_ids=torch.ones(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), ) for patch in ( mock.patch.object(core, "_readout_query_ids", return_value=torch.zeros(1, 1, dtype=torch.long)), mock.patch.object(core.OpenStreamReplayExtractor, "_prepare_pieces", new=prepare), mock.patch.object(core, "_append_open_pieces", side_effect=append), mock.patch.object(core, "_read_open_state", side_effect=read), ): patch.start() self.addCleanup(patch.stop) self.config = core.StreamConfig(memory_seconds=12, readout_last_n_layers=4, readout_layer_fusion="stack") self.kwargs = dict(video_path="unused.mp4", audio_path=None, start=0.0, end=24.0, use_audio_in_video=True, show_progress=False) def multi(self, **kwargs): extractor = MultiWindowOpenStreamReplayExtractor( FakeModel(), object(), self.config, memory_seconds_list=DEFAULT_MEMORY_SECONDS ) return extractor.extract(**{**self.kwargs, **kwargs}) def test_all_caps_match_separate_single_window_extractions_and_encode_once(self): results = self.multi() self.assertEqual(len(self.prepared), 120) self.assertEqual(len(set(self.prepared)), 120) for cap, result in results.items(): reference = core.OpenStreamReplayExtractor( FakeModel(), object(), core.StreamConfig(memory_seconds=cap, readout_last_n_layers=4, readout_layer_fusion="stack") ).extract(**self.kwargs) self.assertTrue(torch.equal(result.times, reference.times)) self.assertTrue(torch.equal(result.states, reference.states), f"mismatch cap={cap}") self.assertEqual(result.states.shape, (120, 4, 3)) self.assertEqual(result.diagnostics["production_processor_calls_shared"], 120) self.assertTrue((result.diagnostics["actual_context_seconds"] <= cap + 1e-6).all()) self.assertTrue(torch.equal(result.diagnostics["state_source_times"], result.times)) self.assertEqual(result.diagnostics["payload_commit_mask"].sum().item(), 12) # Representative physical timeline, independent of any learned scores. index = 102 # cutoff=20.6 s print("Timeline cutoff=20.6s, local A=[19.6,20.6), horizon0.2 target=20.8s") for cap, result in results.items(): start = result.diagnostics["history_start_times"][index].item() self.assertAlmostEqual(start, 22.0 - cap) print(f" cap={cap:g}s global=[{start:g},20.6) age=0s") def test_independent_replay_verification_and_no_future_backfill(self): original = self.multi(verify_every=11) for result in original.values(): self.assertEqual(len(result.verification), 10) self.assertTrue(all(check["max_abs"] == 0 for check in result.verification)) self.future_change = True changed = self.multi() for cap in original: self.assertTrue(torch.equal(original[cap].states[:65], changed[cap].states[:65])) self.assertFalse(torch.equal(original[cap].states[65:], changed[cap].states[65:])) def test_repeated_extract_resets_history_and_supports_mean(self): extractor = MultiWindowOpenStreamReplayExtractor( FakeModel(), object(), core.StreamConfig(), memory_seconds_list=[2, 4] ) first = extractor.extract(**self.kwargs) second = extractor.extract(**self.kwargs) self.assertEqual(first[2].states.shape, (120, 3)) self.assertTrue(torch.equal(first[4].states, second[4].states)) for kwargs in ({"start": 0.2}, {"end": 0.1}, {"end": float("nan")}, {"verify_every": -1}): with self.assertRaises(ValueError): extractor.extract(**{**self.kwargs, **kwargs}) def test_invalid_caps(self): for caps in ([], [2, 2], [2, 2.000000001], [0], [1], [3], [float("nan")], [float("inf")]): with self.subTest(caps=caps), self.assertRaises(ValueError): validate_memory_seconds(caps, 2) def test_export_loads_both_head_and_same_local_audio_and_rejects_bad_alignment(self): results = self.multi(end=4) n = 20 payload = {"mode": "open-stream-replay", "model_id": "synthetic", "start": 0.0, "num_steps": n, "target_path": "synthetic.npy", "config": {"step_seconds": .2}, "label_indices": torch.arange(n), "open_stream_policy": {}} with tempfile.TemporaryDirectory() as directory: root = Path(directory) target_root = root / "targets" target_root.mkdir() targets = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (5, 1)) np.save(target_root / "Session_4.npy", targets) paths = direct.save_multiwindow_banks(payload, results, root / "Session_4.pt") side = {"format": LOCAL_FORMAT, "dataset": "easycom", "sample_id": "Session_4", "partial": False, "times": results[2].times.clone(), "label_indices": torch.arange(n), "audio": torch.ones(n, 5), "audio_valid": torch.ones(n, dtype=torch.bool), "config": {"audio_dim": 5, "causal_policy": "half-open-trailing-window", "uses_speech_labels": False, "encoder_id": "synthetic"}} side_path = root / "audio.pt" torch.save(side, side_path) for cap, path in paths.items(): bank = torch.load(path, weights_only=False) loaded = load_session_data(path.parent, target_root, 4, expected_mode="open-stream-replay") self.assertTrue(torch.equal(loaded.states, results[cap].states)) local, _ = load_local_features(side_path, bank, sample_id="Session_4", dataset="easycom", modalities="audio") self.assertEqual(local.shape, (n, 6)) # encoder vector + availability self.assertFalse(bank["persistent_thinker_kv"]) self.assertEqual(bank["config"]["memory_seconds"], cap) self.assertFalse(bank["partial_extraction"]) for indices in (torch.arange(n) + 1, torch.arange(n - 1)): broken = {**bank, "label_indices": indices} torch.save(broken, path) with self.assertRaises(ValueError): load_session_data(path.parent, target_root, 4, expected_mode="open-stream-replay") with self.assertRaises(ValueError): load_local_features(side_path, broken, sample_id="Session_4", dataset="easycom", modalities="audio") torch.save(bank, path) for broken in ({**bank, "times": bank["times"] + .2}, {**bank, "start": .2}): with self.assertRaises(ValueError): reference_alignment(broken) with self.assertRaises(ValueError): load_session_data(path.parent, target_root, 4) # Wrong mode. with self.assertRaises(FileExistsError): direct.save_multiwindow_banks(payload, results, root / "Session_4.pt") bad = copy.deepcopy(results) bad[12].states[0, 0, 0] = float("nan") with self.assertRaises(ValueError): direct.save_multiwindow_banks(payload, bad, root / "fresh" / "Session_4.pt") self.assertFalse((root / "fresh" / "context_2s" / "Session_4.pt").exists()) class DriverTests(unittest.TestCase): def test_driver_loads_model_once_across_sessions_and_writes_summary(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) target_root = root / "targets" target_root.mkdir() for session in ("Session_4", "Session_5"): (root / session).mkdir() np.save(target_root / f"{session}.npy", np.tile(np.array([[1, 0, 0]], dtype=np.float32), (20, 1))) model, processor = object(), object() def run(command, *, model=None, processor=None): self.assertIs(model, loaded_model) self.assertIs(processor, loaded_processor) args = direct.parse_args(command) times = (torch.arange(20, dtype=torch.float64) + 1) / 5 results = {cap: core.ExtractionResult( times, torch.ones(20, 8), [], {"actual_context_seconds": torch.full((20,), cap)} ) for cap in args.memory_seconds_list} payload = {"mode": "open-stream-replay", "start": 0, "num_steps": 20, "target_path": str(args.target), "config": {"step_seconds": .2}, "label_indices": torch.arange(20), "open_stream_policy": {}} direct.save_multiwindow_banks(payload, results, args.output) loaded_model, loaded_processor = model, processor with mock.patch.object(direct, "_easycom_directory_specs"), \ mock.patch.object(core, "load_qwen_omni", return_value=(model, processor)) as loader, \ mock.patch.object(direct, "main", side_effect=run) as runner: driver.main(["--easycom-video-root", str(root), "--target-root", str(target_root), "--sessions", "Session_4", "Session_5", "--memory-seconds-list", "2", "4", "--output-root", str(root / "output")]) loader.assert_called_once() self.assertEqual(runner.call_count, 2) summary = json.loads((root / "output" / "extraction_summary.json").read_text()) self.assertEqual(set(summary["sessions"]), {"Session_4", "Session_5"}) self.assertFalse(summary["partial_extraction"]) self.assertEqual(summary["sessions"]["Session_5"]["banks"]["4s"]["shape"], [20, 8]) def test_direct_main_routes_multi_banks_and_preserves_scalar_output(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) times = (torch.arange(20, dtype=torch.float64) + 1) / 5 result = core.ExtractionResult(times, torch.ones(20, 8), [], {}) reader = types.SimpleNamespace( clips=[types.SimpleNamespace(path=Path("unused.mp4"))], metadata=lambda: {"synthetic": True} ) model = FakeModel() model.eval = lambda: model model.requires_grad_ = lambda value: model base = ["--video", "unused.mp4", "--end", "4", "--output", str(root / "Session_4.pt")] with mock.patch.object(direct, "CausalSourceAV", return_value=reader), \ mock.patch.object(direct, "_print_media_summary"), \ mock.patch.object(core, "backbone_metadata", return_value={"loader": "unknown"}), \ mock.patch.object(MultiWindowOpenStreamReplayExtractor, "__init__", return_value=None), \ mock.patch.object(MultiWindowOpenStreamReplayExtractor, "extract", return_value={2: result, 4: result}): direct.main(base + ["--memory-seconds-list", "2", "4"], model=model, processor=object()) self.assertFalse((root / "Session_4.pt").exists()) bank = torch.load(root / "context_2s" / "Session_4.pt", weights_only=False) self.assertEqual(bank["config"]["memory_seconds"], 2) self.assertFalse(bank["persistent_thinker_kv"]) self.assertTrue(bank["partial_extraction"]) with mock.patch.object(direct, "CausalSourceAV", return_value=reader), \ mock.patch.object(direct, "_print_media_summary"), \ mock.patch.object(core, "backbone_metadata", return_value={"loader": "unknown"}), \ mock.patch.object(core.OpenStreamReplayExtractor, "__init__", return_value=None), \ mock.patch.object(core.OpenStreamReplayExtractor, "extract", return_value=result): direct.main(base + ["--memory-seconds", "4"], model=model, processor=object()) self.assertTrue((root / "Session_4.pt").exists()) def test_cli_preflight_and_partial_target_policy(self): with tempfile.TemporaryDirectory() as directory, mock.patch.object(direct, "_easycom_directory_specs"): root = Path(directory) (root / "Session_4").mkdir() common = ["--easycom-video-root", str(root), "--sessions", "Session_4", "--output-root", str(root / "output"), "--end", "24"] args = driver.parse_args(common) commands = driver.session_commands(args) self.assertNotIn("--target", commands[0][1]) self.assertIn("--memory-seconds-list", commands[0][1]) self.assertEqual(direct.parse_args(commands[0][1]).readout_last_n_layers, 4) with mock.patch.object(core, "load_qwen_omni") as loader, mock.patch.object(direct, "main") as main: driver.main(common + ["--media-only"]) loader.assert_not_called() main.assert_called_once() for extra in (["--sessions", "Session_99"], ["--sessions", "Session_4", "Session_4"], ["--memory-seconds-list", "3"], ["--end", "0.1"]): with self.assertRaises(ValueError): driver.session_commands(driver.parse_args(common + extra)) existing = root / "output" / "context_2s" / "Session_4.pt" existing.parent.mkdir(parents=True) existing.touch() with self.assertRaises(FileExistsError): driver.session_commands(args) def test_direct_legacy_parser_and_multi_conflicts(self): base = ["--video", "unused.mp4", "--output", "unused.pt", "--end", "24"] self.assertIsNone(direct.parse_args(base).memory_seconds_list) for extra in (["--mode", "causal-prefix"], ["--memory-seconds", "8"], ["--head-checkpoint", "head.pt"], ["--start", ".2"]): with self.assertRaises(SystemExit): direct.main(base + ["--memory-seconds-list", "2", "4"] + extra) if __name__ == "__main__": unittest.main()