Add Qwen tail and chunk ablations
This commit is contained in:
Xin Wang
229
test_streaming_chunk_ablation.py
Normal file
229
test_streaming_chunk_ablation.py
Normal file
@@ -0,0 +1,229 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user