"""Focused tests for the causal EasyCom oracle-information experiment.""" from __future__ import annotations from pathlib import Path from types import SimpleNamespace import tempfile import unittest import numpy as np import torch from easycom_head import TurnPredictionHead from local_turn_features import ORACLE_FORMAT, load_frame_set_features, load_local_features from prepare_easycom_oracle_features import ( causal_history, oracle_person_rows, participant_activity, participant_direct_addressee_activity, ) class OracleFeatureTests(unittest.TestCase): def setUp(self) -> None: torch.set_num_threads(1) def test_transcript_intervals_match_egospeak_right_edge_sampling(self) -> None: rows = [ SimpleNamespace( uid="a", participant_id=3, # EasyCom annotation Start_Frame=1, End_Frame=8. Samples at # source markers 4,8,... therefore include only row 0. dense_start_frame=0, dense_end_frame_exclusive=8, ), SimpleNamespace( uid="b", participant_id=2, # Start_Frame=5, End_Frame=13 includes markers 8 and 12. dense_start_frame=4, dense_end_frame_exclusive=13, ), ] activity, wearer, audit = participant_activity( rows, num_steps=4, wearer_id=2 ) self.assertEqual(activity[3].tolist(), [True, False, False, False]) self.assertEqual(wearer.tolist(), [False, True, True, False]) self.assertEqual( audit["target_sampling_policy"], "egospeak-right-edge-start-inclusive-end-exclusive", ) def test_history_at_row_does_not_read_later_activity(self) -> None: first = torch.tensor([0, 1, 0, 0, 0], dtype=torch.bool) second = first.clone() second[3:] = True observed = causal_history(first, row=2, history_steps=4) changed_future = causal_history(second, row=2, history_steps=4) self.assertTrue(torch.equal(observed, changed_future)) self.assertEqual( observed.tolist(), [0, 0, 1, 0, 0, 1, 1, 1], ) delayed = causal_history(second, row=2, history_steps=4, delay_steps=2) self.assertEqual( delayed.tolist(), [0, 0, 0, 0, 0, 0, 0, 1], ) visibility = torch.tensor([1, 0, 1, 1, 1], dtype=torch.bool) visible_only = causal_history( second, row=2, history_steps=4, observed_sequence=visibility, ) self.assertEqual( visible_only.tolist(), [0, 0, 0, 0, 0, 1, 0, 1], ) def test_direct_addressee_keeps_group_and_unknown_out_of_positive_channel(self) -> None: rows = [ SimpleNamespace( uid="direct", participant_id=3, target_of_speech=(2,), dense_start_frame=0, dense_end_frame_exclusive=8, ), SimpleNamespace( uid="group", participant_id=4, target_of_speech=("Group",), dense_start_frame=4, dense_end_frame_exclusive=13, ), SimpleNamespace( uid="unknown", participant_id=5, target_of_speech=("None",), dense_start_frame=8, dense_end_frame_exclusive=17, ), ] directed, audit = participant_direct_addressee_activity( rows, num_steps=5, wearer_id=2 ) self.assertEqual(directed[3].tolist(), [True, False, False, False, False]) self.assertNotIn(4, directed) self.assertNotIn(5, directed) self.assertEqual( audit["target_of_speech_utterance_counts"], {"direct_wearer": 1, "group": 1, "other": 0, "unknown": 1}, ) self.assertTrue(audit["target_of_speech_group_unknown_not_coerced"]) def test_addressee_oracle_has_capacity_matched_zero_control(self) -> None: activity = {3: torch.tensor([1, 1, 0], dtype=torch.bool)} addressed = {3: torch.tensor([1, 0, 0], dtype=torch.bool)} visibility = {3: torch.ones(3, dtype=torch.bool)} visible = [[3], [3], [3]] records = [ [{"id": "3", "source_clip_index": 0, "source_time": row / 5}] for row in range(3) ] control, control_offsets, _, _ = oracle_person_rows( activity, visibility, visible, records, condition="aggregate-visible-address-control", history_steps=2, delay_steps=0, seed=0, direct_addressee_activity=addressed, ) oracle, oracle_offsets, _, _ = oracle_person_rows( activity, visibility, visible, records, condition="aggregate-visible-addressee", history_steps=2, delay_steps=0, seed=0, direct_addressee_activity=addressed, ) self.assertEqual(control.shape, (3, 6)) self.assertTrue(torch.equal(control_offsets, oracle_offsets)) self.assertTrue(torch.equal(control[:, :4], oracle[:, :4])) self.assertEqual(control[:, 4:].tolist(), [[0, 0], [0, 0], [0, 0]]) self.assertEqual(oracle[:, 4:].tolist(), [[0, 1], [1, 0], [0, 0]]) def test_shuffled_id_preserves_every_lag_count_but_breaks_histories(self) -> None: activity = { 3: torch.tensor([1, 1, 0, 0, 1, 1], dtype=torch.bool), 4: torch.tensor([0, 0, 1, 1, 0, 0], dtype=torch.bool), 5: torch.tensor([0, 1, 0, 1, 0, 1], dtype=torch.bool), } visible = [[3, 4, 5] for _ in range(6)] visibility = { person_id: torch.ones(6, dtype=torch.bool) for person_id in activity } records = [ [ {"id": str(person_id), "source_clip_index": 0, "source_time": row / 5} for person_id in person_ids ] for row, person_ids in enumerate(visible) ] original, offsets, _, _ = oracle_person_rows( activity, visibility, visible, records, condition="per-person", history_steps=4, delay_steps=5, seed=9, ) shuffled, shuffled_offsets, _, audit = oracle_person_rows( activity, visibility, visible, records, condition="shuffled-id", history_steps=4, delay_steps=5, seed=9, ) self.assertTrue(torch.equal(offsets, shuffled_offsets)) for row in range(6): start, stop = int(offsets[row]), int(offsets[row + 1]) self.assertTrue( torch.equal( original[start:stop, :4].sum(dim=0), shuffled[start:stop, :4].sum(dim=0), ) ) self.assertGreater(audit["shuffled_values_changed"], 0) def test_oracle_format_requires_explicit_loader_and_head_mode(self) -> None: n = 4 times = (torch.arange(n, dtype=torch.float64) + 1) / 5 reference = { "states": torch.randn(n, 8), "label_indices": torch.arange(n), "times": times, "mode": "open-stream-replay", "start": 0, } sidecar = { "format": ORACLE_FORMAT, "dataset": "easycom", "sample_id": "Session_4", "times": times.clone(), "label_indices": torch.arange(n), "partial": False, "visual": torch.randn(n, 6), "audio": torch.randn(n, 5), "visual_valid": torch.ones(n, dtype=torch.bool), "audio_valid": torch.ones(n, dtype=torch.bool), "persons": { "features": torch.randn(n, 6), "frame_times": torch.arange(n, dtype=torch.float64) / 5, "frame_offsets": torch.arange(n + 1, dtype=torch.long), "records": [{"id": "3"} for _ in range(n)], }, "config": { "oracle_analysis": True, "uses_speech_labels": True, "causal_policy": "egospeak-right-edge-current-state", "oracle_condition": "per-person", "visual_dim": 6, "audio_dim": 5, }, } with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "Session_4.pt" torch.save(sidecar, path) with self.assertRaisesRegex(ValueError, "format"): load_local_features( path, reference, sample_id="Session_4", dataset="easycom", modalities="av", ) audio, people, contract = load_frame_set_features( path, reference, sample_id="Session_4", dataset="easycom", oracle_analysis=True, ) self.assertEqual(audio.shape, (n, 6)) self.assertEqual(people.dimension, 6) self.assertEqual(contract["format"], ORACLE_FORMAT) self.assertTrue(torch.equal(people.row_to_frame, torch.arange(n))) self.assertTrue(torch.allclose(people.chosen_frame_times, times)) head = TurnPredictionHead( hidden_size=8, num_horizons=2, num_classes=3, local_visual_dim=6, local_audio_dim=5, local_projection_dim=4, local_fusion="oracle-frame-set-attention", person_attention_dim=4, person_attention_heads=1, ) padded, mask = people.padded() logits = head(reference["states"], audio, padded, mask) self.assertEqual(logits.shape, (n, 2, 3)) if __name__ == "__main__": unittest.main()