121 lines
3.6 KiB
Python
121 lines
3.6 KiB
Python
|
"""Focused tests for the direct competing-risk floor-event probe."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from train_easycom_floor_event_probe import (
|
||
|
|
DirectFloorEventHead,
|
||
|
|
FloorEventDataset,
|
||
|
|
ProbeSession,
|
||
|
|
metric_rows,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def activity(wearer: int, others: int) -> list[int]:
|
||
|
|
return [int(not wearer and not others), wearer, others]
|
||
|
|
|
||
|
|
|
||
|
|
def probe_session() -> ProbeSession:
|
||
|
|
target = np.asarray(
|
||
|
|
[
|
||
|
|
activity(0, 1),
|
||
|
|
activity(0, 0),
|
||
|
|
activity(1, 0),
|
||
|
|
activity(1, 0),
|
||
|
|
activity(1, 1),
|
||
|
|
activity(0, 1),
|
||
|
|
activity(0, 1),
|
||
|
|
activity(0, 0),
|
||
|
|
activity(0, 0),
|
||
|
|
activity(0, 1),
|
||
|
|
],
|
||
|
|
dtype=np.float32,
|
||
|
|
)
|
||
|
|
steps = len(target)
|
||
|
|
return ProbeSession(
|
||
|
|
name="Session_12",
|
||
|
|
states=torch.randn(steps, 2, 4),
|
||
|
|
audio=torch.cat((torch.randn(steps, 3), torch.ones(steps, 1)), dim=1),
|
||
|
|
oracle=torch.randn(steps, 6),
|
||
|
|
oracle_valid=torch.ones(steps, dtype=torch.bool),
|
||
|
|
target=target,
|
||
|
|
oracle_condition="aggregate-visible",
|
||
|
|
oracle_contract={"config": {"oracle_condition": "aggregate-visible"}},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class FloorEventProbeTests(unittest.TestCase):
|
||
|
|
def setUp(self) -> None:
|
||
|
|
torch.set_num_threads(1)
|
||
|
|
|
||
|
|
def test_dataset_uses_requested_decision_source(self) -> None:
|
||
|
|
session = probe_session()
|
||
|
|
pre = FloorEventDataset(
|
||
|
|
[session], decision_timing="pre-boundary", midturn_pause_max_steps=10
|
||
|
|
)
|
||
|
|
post = FloorEventDataset(
|
||
|
|
[session], decision_timing="post-boundary", midturn_pause_max_steps=10
|
||
|
|
)
|
||
|
|
pre_response = next(
|
||
|
|
item
|
||
|
|
for item in pre.examples
|
||
|
|
if item.event_index == 0 and item.horizon_index == 1 and item.label
|
||
|
|
)
|
||
|
|
post_response = next(
|
||
|
|
item
|
||
|
|
for item in post.examples
|
||
|
|
if item.event_index == 0 and item.horizon_index == 1 and item.label
|
||
|
|
)
|
||
|
|
self.assertEqual(pre_response.source_row, 0)
|
||
|
|
self.assertEqual(post_response.source_row, 1)
|
||
|
|
|
||
|
|
def test_head_outputs_two_events_by_ten_horizons(self) -> None:
|
||
|
|
model = DirectFloorEventHead(
|
||
|
|
state_shape=(2, 4),
|
||
|
|
audio_dim=3,
|
||
|
|
oracle_dim=6,
|
||
|
|
layer_fusion="concat-mlp",
|
||
|
|
local_projection_dim=4,
|
||
|
|
person_attention_dim=4,
|
||
|
|
person_attention_heads=1,
|
||
|
|
dropout=0.0,
|
||
|
|
)
|
||
|
|
output = model(
|
||
|
|
torch.randn(5, 2, 4),
|
||
|
|
torch.cat((torch.randn(5, 3), torch.ones(5, 1)), dim=1),
|
||
|
|
torch.randn(5, 6),
|
||
|
|
torch.ones(5),
|
||
|
|
)
|
||
|
|
self.assertEqual(output.shape, (5, 2, 10))
|
||
|
|
|
||
|
|
def test_fixed_threshold_is_applied_without_reselection(self) -> None:
|
||
|
|
predictions = {
|
||
|
|
"score": np.tile(np.asarray([0.9, 0.4]), 20),
|
||
|
|
"event": np.repeat(np.arange(2), 20),
|
||
|
|
"horizon": np.tile(np.repeat(np.arange(10), 2), 2),
|
||
|
|
"label": np.tile(np.asarray([True, False]), 20),
|
||
|
|
"short_resume": np.zeros(40, dtype=bool),
|
||
|
|
}
|
||
|
|
thresholds = {
|
||
|
|
(event, horizon): 0.95
|
||
|
|
for event in range(2)
|
||
|
|
for horizon in range(10)
|
||
|
|
}
|
||
|
|
rows, _ = metric_rows(
|
||
|
|
predictions,
|
||
|
|
matched_recall=0.8,
|
||
|
|
fixed_thresholds=thresholds,
|
||
|
|
threshold_source="Session_12",
|
||
|
|
)
|
||
|
|
self.assertEqual(len(rows), 20)
|
||
|
|
self.assertEqual(rows[0]["threshold"], 0.95)
|
||
|
|
self.assertEqual(rows[0]["achieved_positive_recall"], 0.0)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|