197 lines
7.4 KiB
Python
197 lines
7.4 KiB
Python
"""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()
|