Unify FFN and TrajMixer model architectures

evaluate_auc_v2.py

@@ -36,6 +36,7 @@ from delphi2m_auc_report import (
     build_delphi2m_auc_report,
 )
 from eval_data import load_sequence_eval_dataset
+from model_architectures import resolve_model_architecture
 from models import DeepHealth
 from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -184,19 +185,24 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
     return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
 
 
-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,
@@ -207,10 +213,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)),
+        model_architecture=model_architecture,
     )
 
 
-def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
+def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
+    resolve_model_architecture(model.model_architecture, state_dict)
     model.load_state_dict(state_dict, strict=True)
 
 
@@ -1392,12 +1400,17 @@ def main() -> None:
 
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
+    model_architecture = resolve_model_architecture(cfg_model, state_dict)
+    cfg_model["model_architecture"] = model_architecture
+    print(f"Resolved model architecture: {model_architecture}")
 
     device = resolve_eval_device(args.device)
     if device.type == "cuda":
         torch.backends.cudnn.benchmark = True
 
-    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
 
     if (
         model_target_mode == "next_token"