Unify FFN and TrajMixer model architectures

evaluate_auc.py

@@ -45,6 +45,7 @@ from delphi2m_auc_report import (
     build_delphi2m_auc_report,
 )
 from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
+from model_architectures import resolve_model_architecture
 from models import DeepHealth
 from readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
@@ -313,19 +314,24 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
     return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
 
 
-def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
+def build_model_from_dataset(
+    args: argparse.Namespace,
+    cfg: Dict[str, Any],
+    dataset: HealthDataset,
+    state_dict: Optional[Dict[str, Any]] = None,
+) -> DeepHealth:
     model_target_mode = str(cfg_get(
         args, cfg, "model_target_mode", "next_token")).lower()
     if model_target_mode not in {"next_token", "all_future"}:
         raise ValueError(
             f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
         )
+    model_architecture = resolve_model_architecture(cfg, state_dict)
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
         n_head=int(cfg_get(args, cfg, "n_head", 10)),
-        n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
-        n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
+        n_layer=int(cfg["n_layer"]),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -336,6 +342,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)),
+        model_architecture=model_architecture,
     )
 
 
@@ -383,7 +390,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
 
 
 def load_model_state(
-    model: torch.nn.Module,
+    model: DeepHealth,
     checkpoint_path: str,
     device: torch.device,
     state_dict: Optional[Dict[str, Any]] = None,
@@ -391,6 +398,7 @@ def load_model_state(
     state = state_dict if state_dict is not None else load_checkpoint_state_dict(
         checkpoint_path, map_location=device)
 
+    resolve_model_architecture(model.model_architecture, state)
     model.load_state_dict(state, strict=True)
 
 
@@ -1371,14 +1379,19 @@ def main() -> None:
     cfg = dict(cfg)
     cfg["dist_mode"] = dist_mode
     cfg["model_target_mode"] = model_target_mode
+    model_architecture = resolve_model_architecture(cfg, state_dict)
+    cfg["model_architecture"] = model_architecture
     print(f"Resolved dist_mode for evaluation: {dist_mode}")
+    print(f"Resolved model architecture: {model_architecture}")
     print(f"Model target mode for AUC: {model_target_mode}")
     print(
         "AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
         "dist_mode affects model loading but is not converted to horizon-specific risk probability."
     )
 
-    model = build_model_from_dataset(args, cfg, dataset).to(device)
+    model = build_model_from_dataset(
+        args, cfg, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, str(model_ckpt_path),
                      device, state_dict=state_dict)
     model.eval()