import unittest import numpy as np import torch import torch.nn as nn from torch.utils.data import Subset from eval_data import build_model_from_dataset from models import DeepHealth, OtherInfoTokenizer from train_util import fit_continuous_robust_scaler class _ToyAllFutureDataset: def __init__(self): self.cont_type_ids = [1, 3] self.n_types = 4 self.patients = [ self._patient([1, 3], [0.0, 10.0]), self._patient([1, 3], [1.0, 10.0]), self._patient([1, 3], [2.0, 10.0]), self._patient([1, 3], [3.0, 10.0]), self._patient([1, 3], [4.0, 10.0]), self._patient([1, 3], [1000.0, 999.0]), ] @staticmethod def _patient(types, values): return { "other_type": np.asarray(types, dtype=np.int64), "other_value": np.asarray(values, dtype=np.float32), "other_value_kind": np.ones(len(types), dtype=np.int64), } def __len__(self): return len(self.patients) class _CaptureContinuousEncoder(nn.Module): def __init__(self, n_embd): super().__init__() self.n_embd = n_embd self.last_type = None self.last_value = None def forward(self, cont_type_idx, value): self.last_type = cont_type_idx.detach().clone() self.last_value = value.detach().clone() return value[:, None].expand(-1, self.n_embd) class ContinuousValueScalingTests(unittest.TestCase): def test_fit_uses_only_training_subset_and_handles_constant_features(self): dataset = _ToyAllFutureDataset() train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4])) stats = fit_continuous_robust_scaler(dataset, train_subset) self.assertEqual(stats.cont_type_ids, (1, 3)) np.testing.assert_array_equal(stats.observation_count, np.asarray([5, 5])) np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0])) np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0])) def test_fit_supports_next_step_sample_storage(self): dataset = _ToyAllFutureDataset() dataset.samples = dataset.patients del dataset.patients train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4])) stats = fit_continuous_robust_scaler(dataset, train_subset) np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0])) np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0])) def test_tokenizer_standardizes_only_continuous_values(self): tokenizer = OtherInfoTokenizer( n_embd=4, n_types=4, n_cont_types=2, n_categories=3, cont_type_ids=[1, 3], continuous_value_center=[10.0, 100.0], continuous_value_scale=[2.0, 20.0], ) capture = _CaptureContinuousEncoder(n_embd=4) tokenizer.cont_value_encoder = capture tokenizer( other_type=torch.tensor([[1, 2, 3]], dtype=torch.long), other_value=torch.tensor([[14.0, 1.0, 80.0]]), other_value_kind=torch.tensor([[1, 2, 1]], dtype=torch.long), ) torch.testing.assert_close(capture.last_type, torch.tensor([0, 1])) torch.testing.assert_close(capture.last_value, torch.tensor([2.0, -1.0])) def test_scaler_buffers_round_trip_in_new_checkpoint(self): tokenizer = OtherInfoTokenizer( n_embd=4, n_types=4, n_cont_types=2, n_categories=2, cont_type_ids=[1, 3], continuous_value_center=[2.0, 10.0], continuous_value_scale=[1.5, 4.0], ) state = tokenizer.state_dict() self.assertIn("continuous_value_center", state) self.assertIn("continuous_value_scale", state) restored = OtherInfoTokenizer( n_embd=4, n_types=4, n_cont_types=2, n_categories=2, cont_type_ids=[1, 3], continuous_value_center=[0.0, 0.0], continuous_value_scale=[1.0, 1.0], ) restored.load_state_dict(state, strict=True) torch.testing.assert_close( restored.continuous_value_center, torch.tensor([2.0, 10.0]), ) torch.testing.assert_close( restored.continuous_value_scale, torch.tensor([1.5, 4.0]), ) def test_continuous_tokenizer_rejects_missing_scaler_statistics(self): with self.assertRaisesRegex(ValueError, "require train-split RobustScale"): OtherInfoTokenizer( n_embd=4, n_types=4, n_cont_types=2, n_categories=2, cont_type_ids=[1, 3], ) def test_evaluation_rejects_unscaled_continuous_checkpoint(self): dataset = type( "DatasetMetadata", (), { "vocab_size": 8, "n_types": 4, "n_cont_types": 2, "n_categories": 2, "cont_type_ids": [1, 3], }, )() cfg = { "model_target_mode": "all_future", "target_mode": "all_future", "model_architecture": "transformer_ffn_v1", "n_layer": 1, "time_mode": "absolute", "dist_mode": "exponential", } with self.assertRaisesRegex(RuntimeError, "unscaled checkpoints are not supported"): build_model_from_dataset( None, cfg, dataset, state_dict={"blocks.0.mlp.w1.weight": torch.zeros(1)}, ) def test_evaluation_restores_required_scaler_buffers(self): dataset = type( "DatasetMetadata", (), { "vocab_size": 8, "n_types": 4, "n_cont_types": 2, "n_categories": 2, "cont_type_ids": [1, 3], }, )() source = DeepHealth( vocab_size=8, n_embd=4, n_head=1, n_layer=1, n_types=4, n_cont_types=2, n_categories=2, cont_type_ids=[1, 3], continuous_value_center=[2.0, 10.0], continuous_value_scale=[1.5, 4.0], target_mode="all_future", time_mode="absolute", dist_mode="exponential", model_architecture="transformer_ffn_v1", ) state = source.state_dict() cfg = { "model_target_mode": "all_future", "target_mode": "all_future", "model_architecture": "transformer_ffn_v1", "n_embd": 4, "n_head": 1, "n_layer": 1, "n_bins": 16, "time_mode": "absolute", "dist_mode": "exponential", "continuous_value_scaling": "robust", } restored = build_model_from_dataset(None, cfg, dataset, state_dict=state) restored.load_state_dict(state, strict=True) torch.testing.assert_close( restored.tokenizer.continuous_value_center, torch.tensor([2.0, 10.0]), ) torch.testing.assert_close( restored.tokenizer.continuous_value_scale, torch.tensor([1.5, 4.0]), ) if __name__ == "__main__": unittest.main()