Remove legacy event and mixed distribution paths

This commit is contained in:
2026-08-01 14:23:18 +08:00
parent dfb22adf2d
commit de6f9b75b9
22 changed files with 370 additions and 463 deletions

View File

@@ -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()