Implement TrajMixer block

evaluate_auc.py

@@ -40,7 +40,11 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset
 from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
-from models import DeepHealth
+from models import (
+    DeepHealth,
+    validate_traj_mixer_config,
+    validate_traj_mixer_state_dict,
+)
 from readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
 
@@ -309,6 +313,7 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
+    validate_traj_mixer_config(cfg)
     model_target_mode = str(cfg_get(
         args, cfg, "model_target_mode", "next_token")).lower()
     if model_target_mode not in {"next_token", "all_future"}:
@@ -331,6 +336,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
         dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
         dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
+        hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)),
     )
 
 
@@ -386,6 +392,7 @@ def load_model_state(
     state = state_dict if state_dict is not None else load_checkpoint_state_dict(
         checkpoint_path, map_location=device)
 
+    validate_traj_mixer_state_dict(state)
     model.load_state_dict(state, strict=True)