wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add shared-AV multiwindow Qwen feature extraction

Browse Source
This commit is contained in:
Xin Wang
2026-09-17 03:00:05 +08:00
parent a9f748d6a6
commit 9bc89b556f
8 changed files with 971 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

299
test_streaming_multiwindow.py Normal file
View File

@@ -0,0 +1,299 @@
"""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()