Add causal MiniCPM-o streaming feature extraction
This commit is contained in:
Xin Wang
@@ -0,0 +1,230 @@
|
||||
"""CPU-only contract tests for MiniCPM-o streaming feature schedules."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from easycom_head import load_session_data
|
||||
from local_turn_features import reference_alignment
|
||||
from streaming_minicpmo_features_directly import (
|
||||
FEATURE_FORMAT,
|
||||
MiniCPMFeatureExtractor,
|
||||
MiniCPMStreamConfig,
|
||||
_comparison,
|
||||
_time_grid,
|
||||
validate_transformers_version,
|
||||
)
|
||||
|
||||
|
||||
class FakeReader:
|
||||
def __init__(self, duration: float = 3.0) -> None:
|
||||
self.audio = np.arange(
|
||||
int(duration * 16_000), dtype=np.float32
|
||||
) / 16_000.0
|
||||
|
||||
def frame_at_or_before(self, timestamp: float) -> float:
|
||||
return float(timestamp)
|
||||
|
||||
def audio_slice(self, start: float, end: float) -> np.ndarray:
|
||||
first = int(round(start * 16_000))
|
||||
last = int(round(end * 16_000))
|
||||
return self.audio[first:last].copy()
|
||||
|
||||
|
||||
class FakeEngine:
|
||||
hidden_size = 3
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.committed_until = 0.0
|
||||
self.reset_count = 0
|
||||
self.last_restore_audit = None
|
||||
|
||||
def reset(self) -> None:
|
||||
self.committed_until = 0.0
|
||||
self.reset_count += 1
|
||||
|
||||
@staticmethod
|
||||
def _state(frame: float, audio: np.ndarray) -> torch.Tensor:
|
||||
return torch.tensor(
|
||||
[[float(audio[-1]), float(frame), float(len(audio))]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
def commit_state(self, *, frame: float, audio: np.ndarray) -> torch.Tensor:
|
||||
self.committed_until += len(audio) / 16_000.0
|
||||
return self._state(frame, audio)
|
||||
|
||||
def temporary_state(self, *, frame: float, audio: np.ndarray) -> torch.Tensor:
|
||||
before = self.committed_until
|
||||
state = self._state(frame, audio)
|
||||
if self.committed_until != before:
|
||||
raise AssertionError("temporary state modified the fake main cache")
|
||||
self.last_restore_audit = {
|
||||
"kind": "snapshot_restore",
|
||||
"llm_cache_length": 0,
|
||||
"audio_cache_length": 0,
|
||||
"audio_chunk_index": 0,
|
||||
"streaming_position_offset": 0,
|
||||
"window_pruning_suspended": True,
|
||||
"window_policy_restored": True,
|
||||
"llm_checksum_matches": True,
|
||||
"audio_checksum_matches": True,
|
||||
"history_matches": True,
|
||||
"passed": True,
|
||||
}
|
||||
return state
|
||||
|
||||
|
||||
class MiniCPMStreamingTests(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.config = MiniCPMStreamConfig(
|
||||
boundary_max_abs=0.0,
|
||||
boundary_min_cosine=1.0,
|
||||
replay_max_abs=0.0,
|
||||
replay_min_cosine=1.0,
|
||||
)
|
||||
|
||||
def test_time_grid_is_right_edged_and_rejects_fractional_tail(self) -> None:
|
||||
self.assertEqual(_time_grid(0.0, 1.0, 0.2), [0.2, 0.4, 0.6, 0.8, 1.0])
|
||||
with self.assertRaisesRegex(ValueError, "integer number"):
|
||||
_time_grid(0.0, 1.1, 0.2)
|
||||
|
||||
def test_native_hold_has_no_future_backfill(self) -> None:
|
||||
extractor = MiniCPMFeatureExtractor(FakeReader(), FakeEngine(), self.config)
|
||||
result = extractor.extract_native_hold(
|
||||
start=0.0,
|
||||
end=2.0,
|
||||
show_progress=False,
|
||||
)
|
||||
self.assertEqual(tuple(result.states.shape), (10, 3))
|
||||
self.assertTrue(torch.equal(result.states[:4], torch.zeros(4, 3)))
|
||||
self.assertTrue(result.initialization_mask[:4].all())
|
||||
self.assertFalse(result.initialization_mask[4:].any())
|
||||
self.assertEqual(
|
||||
result.native_update_mask.nonzero().flatten().tolist(),
|
||||
[4, 9],
|
||||
)
|
||||
self.assertTrue(torch.equal(result.states[4], result.states[8]))
|
||||
self.assertFalse(torch.equal(result.states[8], result.states[9]))
|
||||
self.assertEqual(
|
||||
result.state_source_times.tolist(),
|
||||
[0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0],
|
||||
)
|
||||
|
||||
def test_cached_partial_updates_every_row_and_passes_audits(self) -> None:
|
||||
reader = FakeReader()
|
||||
extractor = MiniCPMFeatureExtractor(reader, FakeEngine(), self.config)
|
||||
original_audio = reader.audio.copy()
|
||||
result = extractor.extract_cached_partial(
|
||||
start=0.0,
|
||||
end=2.0,
|
||||
verify_every=0,
|
||||
future_intervention_every=0,
|
||||
allow_unverified=False,
|
||||
show_progress=False,
|
||||
)
|
||||
self.assertTrue(result.partial_tail_admissible)
|
||||
self.assertTrue(np.array_equal(reader.audio, original_audio))
|
||||
self.assertFalse(result.initialization_mask.any())
|
||||
self.assertTrue(torch.equal(result.state_source_times, result.times))
|
||||
self.assertEqual(
|
||||
result.native_update_mask.nonzero().flatten().tolist(),
|
||||
[4, 9],
|
||||
)
|
||||
self.assertEqual(
|
||||
{item["kind"] for item in result.verification},
|
||||
{
|
||||
"native_boundary_equivalence",
|
||||
"clean_stream_replay",
|
||||
"future_audio_intervention",
|
||||
"snapshot_restore",
|
||||
},
|
||||
)
|
||||
self.assertTrue(all(item["passed"] for item in result.verification))
|
||||
self.assertFalse(torch.equal(result.states[0], result.states[1]))
|
||||
self.assertLess(float(result.states[0, 1]), 0.2)
|
||||
|
||||
def test_short_partial_run_is_rejected_without_boundary_audit(self) -> None:
|
||||
extractor = MiniCPMFeatureExtractor(FakeReader(), FakeEngine(), self.config)
|
||||
with self.assertRaisesRegex(RuntimeError, "missing=.*native_boundary"):
|
||||
extractor.extract_cached_partial(
|
||||
start=0.0,
|
||||
end=0.8,
|
||||
verify_every=0,
|
||||
future_intervention_every=0,
|
||||
allow_unverified=False,
|
||||
show_progress=False,
|
||||
)
|
||||
|
||||
def test_comparison_rejects_different_states(self) -> None:
|
||||
item = _comparison(
|
||||
torch.zeros(1, 2),
|
||||
torch.ones(1, 2),
|
||||
kind="test",
|
||||
current_time=0.2,
|
||||
max_abs_tolerance=0.0,
|
||||
min_cosine=1.0,
|
||||
)
|
||||
self.assertFalse(item["passed"])
|
||||
self.assertEqual(item["max_abs"], 1.0)
|
||||
|
||||
def test_transformers_version_is_explicitly_gated(self) -> None:
|
||||
self.assertEqual(validate_transformers_version("4.51.0"), "4.51.0")
|
||||
with self.assertRaisesRegex(RuntimeError, "found 4.52.0"):
|
||||
validate_transformers_version("4.52.0")
|
||||
self.assertEqual(
|
||||
validate_transformers_version("4.52.0", allow_untested=True),
|
||||
"4.52.0",
|
||||
)
|
||||
|
||||
def test_saved_payload_loads_in_existing_head_and_unverified_fails(self) -> None:
|
||||
states = torch.randn(20, 8)
|
||||
payload = {
|
||||
"format": FEATURE_FORMAT,
|
||||
"states": states,
|
||||
"times": (torch.arange(20, dtype=torch.float64) + 1) / 5,
|
||||
"label_indices": torch.arange(20),
|
||||
"mode": "causal-prefix",
|
||||
"extractor_mode": "native-1hz-hold",
|
||||
"config": {"step_seconds": 0.2},
|
||||
}
|
||||
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"
|
||||
torch.save(payload, feature_path)
|
||||
np.save(target_root / "Session_4.npy", targets)
|
||||
loaded = load_session_data(feature_root, target_root, 4)
|
||||
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"]))
|
||||
|
||||
bad_reference = dict(payload)
|
||||
bad_reference["times"] = payload["times"] + 0.2
|
||||
with self.assertRaisesRegex(ValueError, "Reference times"):
|
||||
reference_alignment(bad_reference)
|
||||
|
||||
payload["mode"] = "minicpmo-unverified-partial-tail"
|
||||
torch.save(payload, feature_path)
|
||||
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