Implement TrajMixer block

evaluate_auc_v2.py

@@ -29,7 +29,11 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset
 from eval_data import load_sequence_eval_dataset
-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 CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
 
@@ -178,6 +182,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
 
 
 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"}:
@@ -200,10 +205,12 @@ 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)),
     )
 
 
 def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
+    validate_traj_mixer_state_dict(state_dict)
     model.load_state_dict(state_dict, strict=True)