230 lines
8.3 KiB
Python
230 lines
8.3 KiB
Python
|
"""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()
|