300 lines
17 KiB
Python
300 lines
17 KiB
Python
|
"""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()
|