Remove legacy event and mixed distribution paths
This commit is contained in:
@@ -5,7 +5,8 @@ import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import Subset
|
||||
|
||||
from models import OtherInfoTokenizer
|
||||
from eval_data import build_model_from_dataset
|
||||
from models import DeepHealth, OtherInfoTokenizer
|
||||
from train_util import fit_continuous_robust_scaler
|
||||
|
||||
|
||||
@@ -59,6 +60,17 @@ class ContinuousValueScalingTests(unittest.TestCase):
|
||||
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,
|
||||
@@ -66,7 +78,6 @@ class ContinuousValueScalingTests(unittest.TestCase):
|
||||
n_cont_types=2,
|
||||
n_categories=3,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_scaling="robust",
|
||||
continuous_value_center=[10.0, 100.0],
|
||||
continuous_value_scale=[2.0, 20.0],
|
||||
)
|
||||
@@ -89,7 +100,6 @@ class ContinuousValueScalingTests(unittest.TestCase):
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_scaling="robust",
|
||||
continuous_value_center=[2.0, 10.0],
|
||||
continuous_value_scale=[1.5, 4.0],
|
||||
)
|
||||
@@ -103,7 +113,8 @@ class ContinuousValueScalingTests(unittest.TestCase):
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_scaling="robust",
|
||||
continuous_value_center=[0.0, 0.0],
|
||||
continuous_value_scale=[1.0, 1.0],
|
||||
)
|
||||
restored.load_state_dict(state, strict=True)
|
||||
|
||||
@@ -116,27 +127,99 @@ class ContinuousValueScalingTests(unittest.TestCase):
|
||||
torch.tensor([1.5, 4.0]),
|
||||
)
|
||||
|
||||
def test_legacy_none_mode_keeps_old_state_dict_schema(self):
|
||||
tokenizer = OtherInfoTokenizer(
|
||||
n_embd=4,
|
||||
n_types=4,
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
)
|
||||
state = tokenizer.state_dict()
|
||||
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],
|
||||
)
|
||||
|
||||
self.assertNotIn("continuous_value_center", state)
|
||||
self.assertNotIn("continuous_value_scale", state)
|
||||
restored = OtherInfoTokenizer(
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user