"""CPU contract and head tests for no-box full-frame F2/F3.""" from __future__ import annotations import copy from pathlib import Path import tempfile import unittest import numpy as np import torch from easycom_head import TurnPredictionHead, checkpoint_head_config, load_session_data from fullframe_qwen_av import FULLFRAME_AV_FORMAT, load_fullframe_av_features def _fixture(*, partial: bool = False): count = 6 times = (torch.arange(count, dtype=torch.float64) + 1) / 5 reference = { "states": torch.randn(count, 2, 8), "times": times, "label_indices": torch.arange(count), "start": 0, "mode": "open-stream-replay", } sidecar = { "format": FULLFRAME_AV_FORMAT, "dataset": "easycom", "sample_id": "Session_4", "times": times.clone(), "label_indices": torch.arange(count), "partial": partial, "audio": torch.randn(count, 5), "audio_valid": torch.ones(count, dtype=torch.bool), "visual_tokens": { "features": torch.arange(12 * 6, dtype=torch.float32).reshape(12, 6), "frame_times": torch.tensor([0.0, 0.5, 1.0], dtype=torch.float64), "frame_offsets": torch.tensor([0, 4, 8, 12], dtype=torch.long), "grid_shapes": torch.tensor([[2, 2], [2, 2], [2, 2]], dtype=torch.long), "frame_source_clip_indices": torch.zeros(3, dtype=torch.long), }, "config": { "visual_dim": 6, "audio_dim": 5, "visual_fps": 2.0, "causal_policy": "strict-half-open-fullframe-grid", "uses_speech_labels": False, "uses_boxes": False, "visual_feature_kind": "qwen-vision-post-merger-spatial-token", "visual_source": "easycom-full-frame-after-session-banner-crop", "audio_source": "reused-qwen-local-mixed-audio", }, } return reference, sidecar class FullFrameQwenAVTests(unittest.TestCase): def setUp(self): torch.set_num_threads(1) def _load(self, sidecar, reference, *, history=False): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "Session_4.pt" torch.save(sidecar, path) return load_fullframe_av_features( path, reference, sample_id="Session_4", dataset="easycom", history_max_gap_seconds=0.75 if history else None, ) def test_strict_boundary_mapping_and_same_grid_history(self): reference, sidecar = _fixture() audio, frames, contract = self._load(sidecar, reference, history=True) self.assertEqual(audio.shape, (6, 6)) self.assertEqual(frames.row_to_frame.tolist(), [0, 0, 1, 1, 1, 2]) self.assertEqual( frames.chosen_frame_times.tolist(), [0.0, 0.0, 0.5, 0.5, 0.5, 1.0] ) self.assertEqual(contract["format"], FULLFRAME_AV_FORMAT) histories, mask = frames.padded_history( torch.tensor([5]), max_observations=2 ) self.assertEqual(histories.shape, (1, 4, 2, 6)) self.assertTrue(mask.all()) tokens = sidecar["visual_tokens"]["features"] self.assertTrue(torch.equal(histories[0, 0, 0], tokens[4])) self.assertTrue(torch.equal(histories[0, 0, 1], tokens[8])) def test_rejects_partial_and_malformed_contracts(self): reference, sidecar = _fixture(partial=True) with self.assertRaisesRegex(ValueError, "Partial smoke-test"): self._load(sidecar, reference) _, sidecar = _fixture() wrong = copy.deepcopy(sidecar) wrong["config"]["uses_boxes"] = True with self.assertRaisesRegex(ValueError, "causal/no-box"): self._load(wrong, reference) future = copy.deepcopy(sidecar) future["visual_tokens"]["frame_times"][-1] = reference["times"][-1] + 0.01 with self.assertRaisesRegex(ValueError, "fixed causal grid|extends past"): self._load(future, reference) def test_f2_f3_forward_backward_and_checkpoint_config(self): reference, sidecar = _fixture() local, frames, _ = self._load(sidecar, reference, history=True) rows = torch.tensor([0, 2, 3]) states = torch.randn(3, 8) tokens, token_mask = frames.padded(rows) f2 = TurnPredictionHead( hidden_size=8, local_visual_dim=6, local_audio_dim=5, local_projection_dim=4, local_fusion="fullframe-spatial-attention", person_attention_dim=4, person_attention_heads=2, person_history_frames=2, ) logits = f2(states, local[rows], tokens, token_mask) self.assertEqual(logits.shape, (3, 10, 3)) logits.sum().backward() self.assertIsNotNone(f2.person_attention.in_proj_weight.grad) histories, history_mask = frames.padded_history(rows, max_observations=2) f3 = TurnPredictionHead( hidden_size=8, local_visual_dim=6, local_audio_dim=5, local_projection_dim=4, local_fusion="fullframe-spatiotemporal-attention", person_attention_dim=4, person_attention_heads=2, person_history_frames=2, ) logits = f3(states, local[rows], histories, history_mask) self.assertEqual(logits.shape, (3, 10, 3)) logits.sum().backward() self.assertIsNotNone(f3.person_track_gru.weight_ih_l0.grad) config = checkpoint_head_config( { "head_config": { "hidden_size": 8, "local_visual_dim": 6, "local_audio_dim": 5, "local_fusion": "fullframe-spatiotemporal-attention", "person_history_frames": 2, } } ) self.assertEqual(config["local_fusion"], "fullframe-spatiotemporal-attention") self.assertEqual(config["person_history_frames"], 2) def test_session_loader_routes_fullframe_contract(self): reference, sidecar = _fixture() labels = np.tile( np.asarray( [[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32 ), (2, 1), ) with tempfile.TemporaryDirectory() as directory: root = Path(directory) feature_root = root / "global" target_root = root / "target" local_root = root / "fullframe" for folder in (feature_root, target_root, local_root): folder.mkdir() torch.save(reference, feature_root / "Session_4.pt") torch.save(sidecar, local_root / "Session_4.pt") np.save(target_root / "Session_4.npy", labels) data = load_session_data( feature_root, target_root, "Session_4", expected_mode="open-stream-replay", local_feature_root=local_root, local_modalities="av", local_fusion="fullframe-spatiotemporal-attention", person_track_max_gap_seconds=0.75, ) self.assertEqual(data.local_features.shape, (6, 6)) self.assertIsNotNone(data.local_person_features.previous_observation) self.assertEqual( data.feature_metadata["local_contract"]["format"], FULLFRAME_AV_FORMAT, ) if __name__ == "__main__": unittest.main()