Enhance data parsing and validation, add extra info types files

- Improved `parse_int_list` and `parse_float_list` functions to support JSON list input.
- Introduced `validate_dataset_metadata` function to ensure dataset metadata consistency with training configuration.
- Added multiple new files for extra information types, categorizing them into assessment-only, exposure-only, and combined types.
- Removed deprecated `merge_extra_info_types` function and adjusted related logic in `train.py`.
- Updated `save_config` function to accept additional metadata for training runs.
- Refactored model and training scripts for better clarity and maintainability.

train.py

@@ -20,7 +20,7 @@ import sys
 import time
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Dict, Iterable, Tuple
+from typing import Any, Dict, Tuple
 
 import numpy as np
 import torch
@@ -135,22 +135,6 @@ def load_extra_info_types_file(path: str) -> list[int]:
     return parsed
 
 
-def merge_extra_info_types(*sources: Iterable[int] | None) -> list[int] | None:
-    """Merge optional type-id lists while preserving first-seen order."""
-    merged: list[int] = []
-    seen: set[int] = set()
-    for source in sources:
-        if source is None:
-            continue
-        for raw_type in source:
-            type_id = int(raw_type)
-            if type_id in seen:
-                continue
-            seen.add(type_id)
-            merged.append(type_id)
-    return merged or None
-
-
 def configure_torch_for_training(device: torch.device) -> None:
     """Enable backend settings that can improve training throughput on CUDA."""
     if device.type == "cuda":
@@ -573,19 +557,68 @@ def save_checkpoint(
     torch.save(model.state_dict(), checkpoint_path)
 
 
-def save_config(args: argparse.Namespace, config_path: Path) -> None:
+def save_config(
+    args: argparse.Namespace,
+    config_path: Path,
+    extra: Dict[str, Any] | None = None,
+) -> None:
     """Save training config as JSON."""
-    config_dict = vars(args)
-    # Convert non-serializable types
-    config_dict = {
-        k: str(v) if not isinstance(
-            v, (int, float, str, bool, type(None))) else v
-        for k, v in config_dict.items()
-    }
+    config_dict: Dict[str, Any] = {}
+    for key, value in vars(args).items():
+        if isinstance(value, tuple):
+            config_dict[key] = list(value)
+        elif isinstance(value, list):
+            config_dict[key] = value
+        elif isinstance(value, (int, float, str, bool, type(None))):
+            config_dict[key] = value
+        else:
+            config_dict[key] = str(value)
+    if extra:
+        config_dict.update(extra)
     with open(config_path, "w") as f:
         json.dump(config_dict, f, indent=2)
 
 
+def build_run_metadata(
+    args: argparse.Namespace,
+    dataset: HealthDataset,
+    train_subset: Subset,
+    val_subset: Subset,
+    test_subset: Subset,
+    run_name: str,
+) -> Dict[str, Any]:
+    """Collect resolved training facts needed to rebuild the model for evaluation."""
+    return {
+        "run_name": run_name,
+        "dataset_class": "NextStepHealthDataset",
+        "collate_fn": "next_step_collate_fn",
+        "model_class": "DeepHealth",
+        "model_target_mode": "next_token",
+        "dist_mode": "exponential",
+        "extra_info_types_file": (
+            Path(args.extra_info_types_file).name
+            if args.extra_info_types_file is not None
+            else None
+        ),
+        "extra_info_types": dataset.extra_info_types,
+        "dataset_metadata": {
+            "vocab_size": int(dataset.vocab_size),
+            "n_types": int(dataset.n_types),
+            "n_cont_types": int(dataset.n_cont_types),
+            "n_categories": int(dataset.n_categories),
+            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
+            "extra_info_types": [int(x) for x in dataset.extra_info_types],
+        },
+        "split_sizes": {
+            "train": int(len(train_subset)),
+            "val": int(len(val_subset)),
+            "test": int(len(test_subset)),
+        },
+        "resolved_readout_name": args.readout_name,
+        "resolved_loss_name": args.loss_name,
+    }
+
+
 def normalize_training_config(args: argparse.Namespace) -> None:
     """Fill in and validate training options that depend on other flags."""
     if args.target_mode not in {"delphi2m", "uts"}:
@@ -661,8 +694,6 @@ def main():
                         help="Time encoding mode for disease history")
     parser.add_argument("--dropout", type=float, default=0.0,
                         help="Dropout rate")
-    parser.add_argument("--extra_info_types", type=int, nargs="*", default=None,
-                        help="Optional list of other-information type ids to include")
     parser.add_argument("--extra_info_types_file", type=str, default=None,
                         help="Optional file containing other-information type ids to include")
 
@@ -715,15 +746,11 @@ def main():
                         help="Device to use for training: cpu, cuda, or cuda:<index>")
 
     args = parser.parse_args()
-    file_extra_info_types = (
+    args.extra_info_types = (
         load_extra_info_types_file(args.extra_info_types_file)
         if args.extra_info_types_file is not None
         else None
     )
-    args.extra_info_types = merge_extra_info_types(
-        args.extra_info_types,
-        file_extra_info_types,
-    )
 
     # ---- Setup ----
     set_seed(args.seed)
@@ -893,7 +920,18 @@ def main():
         raise ValueError(f"Unknown loss: {args.loss_name}")
 
     # ---- Save Config ----
-    save_config(args, run_dir / "train_config.json")
+    save_config(
+        args,
+        run_dir / "train_config.json",
+        extra=build_run_metadata(
+            args=args,
+            dataset=dataset,
+            train_subset=train_subset,
+            val_subset=val_subset,
+            test_subset=test_subset,
+            run_name=run_name,
+        ),
+    )
     logger.info(f"Config saved to {run_dir / 'train_config.json'}")
 
     # ---- Training Loop ----