Unify FFN and TrajMixer model architectures

run_missing_training_runs.sh

@@ -10,7 +10,8 @@ set -euo pipefail
 #   all_future + relative time + mixed death/risk head
 #
 # This script only launches those missing training jobs. It intentionally does
-# not call evaluate_*.py and does not add extra random seeds.
+# not call evaluate_*.py and does not add extra random seeds. Set
+# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
 
@@ -18,6 +19,8 @@ PYTHON_BIN="${PYTHON_BIN:-python}"
 DEVICE="${DEVICE:-cuda}"
 NUM_WORKERS="${NUM_WORKERS:-4}"
 PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
+MODEL_ARCHITECTURE="${MODEL_ARCHITECTURE:-transformer_ffn_v1}"
+N_LAYER="${N_LAYER:-12}"
 
 TIME_MODE="relative"
 DIST_MODE="mixed"
@@ -36,8 +39,8 @@ COMMON_ARGS=(
   --min_future_events 1
   --n_embd 120
   --n_head 10
-  --n_hist_layer 12
-  --n_tab_layer 4
+  --n_layer "${N_LAYER}"
+  --model_architecture "${MODEL_ARCHITECTURE}"
   --n_bins 16
   --extra_pool_reduce mean
   --dropout 0.0
@@ -57,15 +60,23 @@ COMMON_ARGS=(
 
 already_trained() {
   local extra_file="$1"
-  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" <<'PY'
+  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
 import json
 import sys
 from pathlib import Path
 
-time_mode, dist_mode, extra_file, seed, validation_query_seed = sys.argv[1:6]
+(
+    time_mode,
+    dist_mode,
+    extra_file,
+    seed,
+    validation_query_seed,
+    model_architecture,
+    n_layer,
+) = sys.argv[1:8]
 extra_name = Path(extra_file).name
 
-for config_path in Path("runs").glob("*/train_config.json"):
+for config_path in Path("runs").rglob("train_config.json"):
     try:
         cfg = json.loads(config_path.read_text(encoding="utf-8"))
     except Exception:
@@ -78,6 +89,8 @@ for config_path in Path("runs").glob("*/train_config.json"):
 
     if (
         cfg.get("model_target_mode") == "all_future"
+        and cfg.get("model_architecture") == model_architecture
+        and int(cfg.get("n_layer", -1)) == int(n_layer)
         and cfg.get("time_mode") == time_mode
         and cfg.get("dist_mode") == dist_mode
         and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
@@ -100,7 +113,7 @@ train_if_missing() {
     return 2
   fi
 
-  echo "==> Checking ${label}: ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
+  echo "==> Checking ${label}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
   if existing_run="$(already_trained "$extra_file")"; then
     echo "    skip: already trained at ${existing_run}"
     return 0