Enhance training script and evaluation logic to support all-future model target mode and improve error handling for distribution modes

train.py

@@ -733,6 +733,11 @@ def normalize_training_config(args: argparse.Namespace) -> None:
         raise ValueError(f"Unknown model_target_mode: {args.model_target_mode}")
     if args.dist_mode not in {"exponential", "weibull", "mixed"}:
         raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
+    if args.model_target_mode == "next_token" and args.dist_mode != "exponential":
+        raise ValueError(
+            "next_token training currently supports dist_mode='exponential' only. "
+            "Use model_target_mode='all_future' for weibull or mixed distributions."
+        )
     if args.all_future_min_history_events < 1:
         raise ValueError("all_future_min_history_events must be >= 1")
     if args.all_future_min_future_events < 1:
@@ -833,7 +838,7 @@ def main():
                         help="Model forward/training mode")
     parser.add_argument("--dist_mode", type=str, default="exponential",
                         choices=["exponential", "weibull", "mixed"],
-                        help="Event-time distribution for model heads and all-future loss")
+                        help="Event-time distribution. next_token requires exponential; all_future supports exponential, weibull, and mixed")
     parser.add_argument("--dropout", type=float, default=0.0,
                         help="Dropout rate")
     parser.add_argument("--extra_info_types_file", type=str, default=None,