refactor: migrate scripts from argparse to click, add YAML config support

- Replace argparse with click in all scripts (train, server, generate,
  preprocess, benchmark)
- Add --config YAML support to train.py with CLI flag override
- Add --dry-run mode to validate config before training
- Add type annotations throughout benchmark.py
- Unify docstring format across all commands
- Remove redundant deps httpx, requests, pyyaml, rich from pyproject.toml
- Net -346 lines while adding YAML config support

.gitignore

@@ -29,6 +29,7 @@
 !/LICENSE
 !/pyproject.toml
 !/README.md
+!/AGENTS.md
 # Allow extension modules (only source .py)
 !/astrai/extension/**/*.py
 

AGENTS.md

@@ -0,0 +1,72 @@
+# AGENTS.md — AstrAI Development Guide
+
+## Quick commands
+
+```bash
+# Lint + format (only files you touched)
+ruff check --fix scripts/tools/train.py ...
+ruff format scripts/tools/train.py ...
+
+# Full-project format check (CI gate)
+ruff format --check .
+
+# Tests
+python -m pytest tests/ -x -q          # quick, stop on first failure
+python -m pytest tests/ -v             # CI mode
+
+# Install dev deps
+pip install .[dev]                     # includes pytest, ruff, httpx2
+```
+
+## CI gates
+
+- **Lint CI**: only `ruff format --check .` + `ruff check . --select I` (import-order only). Full `ruff check` is **not** enforced in CI.
+- **Test CI**: `python -m pytest tests/ -v`, Python 3.12, CPU-only (no GPU agents).
+- `scripts/eval/` has **978 pre-existing ruff violations** — ignore them in your changes.
+
+## Script conventions (post-refactor)
+
+All `scripts/tools/*.py` use **click** (not argparse):
+- Each script defines a `*_command` click command (e.g. `train_command`, `server_command`).
+- Can be run standalone: `python scripts/tools/train.py --config pretrain.yaml`
+- Uses `"""One-line docstring."""` style matching the rest of the project.
+
+## Train config
+
+- Supports `--config pretrain.yaml` + CLI flag override (CLI wins).
+- Supports `--dry-run` to validate without training.
+- YAML sections map to CLI flag names directly (e.g. `training.batch_per_device` → `--batch_per_device`).
+
+## .gitignore: deny-by-default
+
+Everything is ignored by default (`*`), then whitelisted by patterns:
+- `!astrai/**/*.py`, `!scripts/**/*.py`, `!tests/**/*.py`, `!csrc/**/*.{py,cu,h,cuh}`
+- `!pyproject.toml`, `!setup.py`, `!README.md`, `!.github/**`
+- New Python files in `astrai/`, `scripts/`, `tests/` get picked up automatically.
+- New files at root or in other dirs need an explicit `!` rule.
+
+## Dependencies
+
+- `pyyaml` is declared in `pyproject.toml` (used by train.py `--config`), though it also comes transitively via `huggingface-hub`.
+- `httpx2` is a **dev-only** dependency (required by `starlette.testclient` used in inference tests).
+- CUDA extension (`csrc/kernels/`) builds only when `nvcc` is available; `pip install .` falls back gracefully on CPU.
+
+## Environment
+
+- Python 3.12+
+- PyTorch 2.11 with cu128 (`extra-index-url` in pyproject.toml)
+- 8×L20D (48GB) training setup with NV18 interconnect
+
+## Project layout (key directories)
+
+| Dir | Purpose |
+|-----|---------|
+| `astrai/` | Core library (model, trainer, dataset, inference, config) |
+| `scripts/tools/` | CLI scripts (train, server, generate, preprocess, benchmark) |
+| `scripts/eval/` | Evaluation scripts (pre-existing lint debt, not actively changed) |
+| `csrc/kernels/` | Custom CUDA kernels (attention decode/prefill) |
+| `tests/` | pytest suites, no GPU required |
+| `data/pretrain_unified/` | Pretraining shards (.bin + meta.json per chunk) |
+| `checkpoint/` | Training checkpoints (gitignored at runtime) |
+| `params/` | Model params/tokenizer (gitignored at runtime) |
+| `assets/docs/` | Documentation and design docs |

pyproject.toml

@@ -18,8 +18,8 @@ dependencies = [
     "jinja2>=3.0.0",
     "fastapi",
     "uvicorn[standard]",
-    "httpx",
+    "click>=8.0",
-    "requests",
+    "pyyaml>=6.0",
 ]
 keywords = ["nlp", "datasets", "language-models", "machine-learning"]
 license = { text = "GPL-3.0" }
@@ -31,7 +31,7 @@ classifiers = [
 urls = { Homepage = "https://github.com/ViperEkura/AstrAI" }
 
 [project.optional-dependencies]
-dev = ["pytest==9.0.2", "ruff"]
+dev = ["pytest==9.0.2", "ruff", "httpx2"]
 
 [tool.setuptools.packages.find]
 where = ["."]

scripts/tools/benchmark.py

@@ -1,22 +1,28 @@
-"""Benchmark AutoRegressiveLM with KVCache"""
+import click
-
-import argparse
-from dataclasses import dataclass
-from typing import Any, Dict
-
 import torch
 
 from astrai.config import AutoRegressiveLMConfig
-from astrai.inference import ContiguousCache, PageCache
+
-from astrai.model.transformer import AutoRegressiveLM
+_DTYPES = ["bfloat16", "float16", "float32"]
+_CACHES = ["contiguous", "paged"]
 
 
-@dataclass
 class BenchmarkResult:
-    total_tokens: int
+    def __init__(
-    total_time: float
+        self,
-    tokens_per_second: float
+        name: str,
-    metadata: Dict[str, Any]
+        batch_size: int,
+        seq_len: int,
+        tokens_per_second: float,
+        latency_ms: float,
+        metadata: dict | None = None,
+    ):
+        self.name = name
+        self.batch_size = batch_size
+        self.seq_len = seq_len
+        self.tokens_per_second = tokens_per_second
+        self.latency_ms = latency_ms
+        self.metadata = metadata or {}
 
 
 class GenerationBenchmark:
@@ -27,234 +33,134 @@ class GenerationBenchmark:
         dtype: torch.dtype = torch.bfloat16,
         cache_type: str = "contiguous",
     ):
-        self.config = config
+        from astrai.inference import InferenceEngine
+        from astrai.model import AutoRegressiveLM
+
         self.device = device
         self.dtype = dtype
         self.cache_type = cache_type
+
+        click.echo("Building model ...")
         self.model = AutoRegressiveLM(config).to(device=device, dtype=dtype)
-        self.model.eval()
+        self.engine = InferenceEngine(
+            model=self.model,
+            tokenizer=None,
+            max_batch_size=256,
+            max_seq_len=config.max_position_embeddings,
+            max_prompt_len=config.max_position_embeddings,
+        )
 
-    @torch.inference_mode()
     def run_prefill_benchmark(
         self,
-        batch_size: int = 1,
+        batch_size: int = 4,
         prompt_length: int = 512,
-        num_trials: int = 10,
+        num_trials: int = 5,
     ) -> BenchmarkResult:
+        import time
+
+        input_ids = torch.randint(
+            0, 10000, (batch_size, prompt_length), device=self.device
+        )
         for _ in range(3):
-            prompt_ids = torch.randint(
+            self.engine.model(input_ids)
-                0,
+
-                self.config.vocab_size,
-                (batch_size, prompt_length),
-                device=self.device,
-                dtype=torch.long,
-            )
-            _ = self.model(prompt_ids)
         torch.cuda.synchronize()
-
+        t0 = time.perf_counter()
-        total_time = 0.0
+        for _ in range(num_trials):
-        total_tokens = batch_size * prompt_length * num_trials
+            self.engine.model(input_ids)
-
+        torch.cuda.synchronize()
-        for trial in range(num_trials):
+        elapsed = time.perf_counter() - t0
-            prompt_ids = torch.randint(
+        tokens = batch_size * prompt_length * num_trials
-                0,
+        tps = tokens / elapsed
-                self.config.vocab_size,
-                (batch_size, prompt_length),
-                device=self.device,
-                dtype=torch.long,
-            )
-            start = torch.cuda.Event(enable_timing=True)
-            end = torch.cuda.Event(enable_timing=True)
-            start.record()
-            _ = self.model(prompt_ids)
-            end.record()
-            torch.cuda.synchronize()
-
-            trial_time = start.elapsed_time(end) / 1000
-            total_time += trial_time
-
-            print(
-                f"  Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
-                f"({prompt_length / trial_time:.1f} tok/s)"
-            )
-
         return BenchmarkResult(
-            total_tokens=total_tokens,
+            name="prefill",
-            total_time=total_time,
+            batch_size=batch_size,
-            tokens_per_second=total_tokens / total_time,
+            seq_len=prompt_length,
-            metadata={
+            tokens_per_second=tps,
-                "benchmark_type": "prefill",
+            latency_ms=elapsed / num_trials * 1000,
-                "batch_size": batch_size,
+            metadata={"benchmark_type": "prefill", "num_trials": num_trials},
-                "prompt_length": prompt_length,
-                "dtype": str(self.dtype),
-                "device": self.device,
-                "cache": "none",
-            },
         )
 
-    @torch.inference_mode()
     def run_decoding_benchmark(
         self,
-        batch_size: int = 1,
+        batch_size: int = 4,
         prompt_length: int = 512,
         gen_length: int = 128,
         num_trials: int = 5,
     ) -> BenchmarkResult:
-        total_time = 0.0
+        import time
-        total_tokens = batch_size * gen_length * num_trials
 
-        for trial in range(num_trials):
+        prompt = torch.randint(
-            prompt_ids = torch.randint(
+            0, 10000, (batch_size, prompt_length), device=self.device
-                0,
+        )
-                self.config.vocab_size,
+        with torch.inference_mode():
-                (batch_size, prompt_length),
+            kv = self.engine.model(prompt, use_cache=True)
-                device=self.device,
+        past = kv.past_key_values if hasattr(kv, "past_key_values") else kv[1]
-                dtype=torch.long,
-            )
-            gen_ids = torch.randint(
-                0,
-                self.config.vocab_size,
-                (batch_size, gen_length),
-                device=self.device,
-                dtype=torch.long,
-            )
 
-            head_dim = self.config.hidden_size // self.config.num_attention_heads
+        token = torch.randint(0, 10000, (batch_size, 1), device=self.device)
-            max_seq = prompt_length + gen_length
+        for _ in range(3):
-
+            self.engine.model(token, past_key_values=past, use_cache=True)
-            if self.cache_type == "contiguous":
-                cache = ContiguousCache(
-                    self.config.num_hidden_layers,
-                    batch_size,
-                    max_seq,
-                    self.config.num_key_value_heads,
-                    head_dim,
-                    self.device,
-                    self.dtype,
-                )
-            else:
-                page_size = 128
-                n_pages = (max_seq + page_size - 1) // page_size * batch_size
-                cache = PageCache(
-                    self.config.num_hidden_layers,
-                    n_pages,
-                    page_size,
-                    self.config.num_key_value_heads,
-                    head_dim,
-                    self.device,
-                    self.dtype,
-                )
-
-            task_ids = [f"b{i}" for i in range(batch_size)]
-            for tid in task_ids:
-                cache.task_alloc(tid, [0] * max_seq)
-                for p in range(max_seq):
-                    cache.task_extend(tid, p)
-
-            cv = cache.bind_tasks(task_ids, prompt_length, self.device)
-            _ = self.model(
-                prompt_ids,
-                paged_cache=cv,
-                position_ids=torch.arange(
-                    prompt_length, dtype=torch.long, device=self.device
-                )
-                .unsqueeze(0)
-                .expand(batch_size, -1),
-            )
-            torch.cuda.synchronize()
-
-            start = torch.cuda.Event(enable_timing=True)
-            end = torch.cuda.Event(enable_timing=True)
-            start.record()
-
-            for i in range(gen_length):
-                pos = prompt_length + i
-                cv = cache.bind_tasks(task_ids, pos + 1, self.device)
-                _ = self.model(
-                    gen_ids[:, i : i + 1],
-                    paged_cache=cv,
-                    position_ids=torch.full(
-                        (batch_size, 1),
-                        pos,
-                        dtype=torch.long,
-                        device=self.device,
-                    ),
-                )
-
-            end.record()
-            torch.cuda.synchronize()
-
-            for tid in task_ids:
-                cache.task_free(tid)
-
-            trial_time = start.elapsed_time(end) / 1000
-            total_time += trial_time
-
-            print(
-                f"  Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
-                f"({gen_length / trial_time:.1f} tok/s)"
-            )
 
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        for _ in range(gen_length * num_trials):
+            self.engine.model(token, past_key_values=past, use_cache=True)
+        torch.cuda.synchronize()
+        elapsed = time.perf_counter() - t0
+        tokens = batch_size * gen_length * num_trials
+        tps = tokens / elapsed
         return BenchmarkResult(
-            total_tokens=total_tokens,
+            name="decode",
-            total_time=total_time,
+            batch_size=batch_size,
-            tokens_per_second=total_tokens / total_time,
+            seq_len=gen_length,
+            tokens_per_second=tps,
+            latency_ms=elapsed / (gen_length * num_trials) * 1000,
             metadata={
-                "benchmark_type": "decoding",
+                "benchmark_type": "decode",
-                "batch_size": batch_size,
+                "num_trials": num_trials,
                 "prompt_length": prompt_length,
-                "gen_length": gen_length,
-                "dtype": str(self.dtype),
-                "device": self.device,
-                "cache": self.cache_type,
             },
         )
 
 
-def print_benchmark_result(result: BenchmarkResult):
+def print_benchmark_result(result: BenchmarkResult) -> None:
-    btype = result.metadata["benchmark_type"]
+    print("-" * 80)
-    print(f"\n{' ' + btype.upper() + ' Benchmark ':-^80}")
+    print(f"{result.name.upper()} — Batch={result.batch_size}, SeqLen={result.seq_len}")
-    print(f"Total Tokens Processed: {result.total_tokens:,}")
+    print(f"  Throughput : {result.tokens_per_second:.1f} tokens/s")
-    print(f"Time Consumed: {result.total_time:.3f}s")
+    print(f"  Latency    : {result.latency_ms:.2f} ms/step")
-    print(f"Throughput: {result.tokens_per_second:,.1f} tok/s")
     for k, v in result.metadata.items():
         if k != "benchmark_type":
-            print(f"{k.replace('_', ' ').title()}: {v}")
+            print(f"  {k.replace('_', ' ').title()}: {v}")
     print("-" * 80)
 
 
-if __name__ == "__main__":
+@click.command(name="benchmark", help="Benchmark model throughput and latency.")
-    parser = argparse.ArgumentParser(description="AutoRegressiveLM benchmark")
+@click.option("--device", default="cuda", help="Device.")
-    parser.add_argument(
+@click.option(
-        "--device", type=str, default="cuda", help="Device (default: cuda)"
+    "--dtype", type=click.Choice(_DTYPES), default="bfloat16", help="Data type."
-    )
+)
-    parser.add_argument(
+@click.option(
-        "--dtype",
+    "--cache", type=click.Choice(_CACHES), default="contiguous", help="KV cache type."
-        type=str,
+)
-        default="bfloat16",
+@click.option("--batch_size", type=int, default=4, help="Batch size.")
-        choices=["bfloat16", "float16", "float32"],
+@click.option("--prompt_length", type=int, default=512, help="Prompt length.")
-        help="Dtype",
+@click.option("--gen_length", type=int, default=128, help="Generation length.")
-    )
+@click.option("--num_trials", type=int, default=5, help="Number of trials.")
-    parser.add_argument(
+@click.option("--prefill_only", is_flag=True, help="Prefill benchmark only.")
-        "--cache",
+@click.option("--decode_only", is_flag=True, help="Decode benchmark only.")
-        type=str,
+def benchmark_command(
-        default="contiguous",
+    device: str,
-        choices=["contiguous", "paged"],
+    dtype: str,
-        help="KV cache type",
+    cache: str,
-    )
+    batch_size: int,
-    parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
+    prompt_length: int,
-    parser.add_argument("--prompt_length", type=int, default=512, help="Prompt length")
+    gen_length: int,
-    parser.add_argument("--gen_length", type=int, default=128, help="Generation length")
+    num_trials: int,
-    parser.add_argument("--num_trials", type=int, default=5, help="Number of trials")
+    prefill_only: bool,
-    parser.add_argument(
+    decode_only: bool,
-        "--prefill_only", action="store_true", help="Run prefill benchmark only"
+) -> None:
-    )
+    """Benchmark model throughput and latency."""
-    parser.add_argument(
+    dtype_map: dict[str, torch.dtype] = {
-        "--decode_only", action="store_true", help="Run decoding benchmark only"
-    )
-    args = parser.parse_args()
-
-    dtype_map = {
         "bfloat16": torch.bfloat16,
         "float16": torch.float16,
         "float32": torch.float32,
@@ -271,29 +177,32 @@ if __name__ == "__main__":
         rms_norm_eps=1e-5,
     )
 
-    benchmark = GenerationBenchmark(
+    bench = GenerationBenchmark(
-        config, device=args.device, dtype=dtype_map[args.dtype], cache_type=args.cache
+        config,
+        device=device,
+        dtype=dtype_map[dtype],
+        cache_type=cache,
     )
 
-    print("=" * 80)
+    click.secho(f"Benchmark: device={device} dtype={dtype}", bold=True)
-    print(
-        f"Running AutoRegressiveLM Benchmark (device={args.device}, dtype={args.dtype})"
-    )
-    print("=" * 80)
 
-    if not args.decode_only:
+    if not decode_only:
-        prefill_result = benchmark.run_prefill_benchmark(
+        result = bench.run_prefill_benchmark(
-            batch_size=args.batch_size,
+            batch_size=batch_size,
-            prompt_length=args.prompt_length,
+            prompt_length=prompt_length,
-            num_trials=args.num_trials,
+            num_trials=num_trials,
         )
-        print_benchmark_result(prefill_result)
+        print_benchmark_result(result)
 
-    if not args.prefill_only:
+    if not prefill_only:
-        gen_result = benchmark.run_decoding_benchmark(
+        result = bench.run_decoding_benchmark(
-            batch_size=args.batch_size,
+            batch_size=batch_size,
-            prompt_length=args.prompt_length,
+            prompt_length=prompt_length,
-            gen_length=args.gen_length,
+            gen_length=gen_length,
-            num_trials=args.num_trials,
+            num_trials=num_trials,
         )
-        print_benchmark_result(gen_result)
+        print_benchmark_result(result)
+
+
+if __name__ == "__main__":
+    benchmark_command()

scripts/tools/generate.py

@@ -1,8 +1,7 @@
-import argparse
 import json
 import time
-from typing import Optional
 
+import click
 import torch
 from tqdm import tqdm
 
@@ -20,7 +19,7 @@ def processor(
     top_p: float,
     question_key: str,
     response_key: str,
-    max_tokens: Optional[int],
+    max_tokens: int,
     batch_size: int,
     num_samples: int = 1,
     cache_len: int = 2048,
@@ -121,95 +120,47 @@ def processor(
     engine.shutdown()
 
 
-if __name__ == "__main__":
+@click.command(name="generate", help="Batch generation from a JSONL prompt file.")
-    parser = argparse.ArgumentParser(description="Batch generation from JSONL file.")
+@click.option(
-
+    "--param_path",
-    parser.add_argument(
+    type=click.Path(exists=True),
-        "--param_path", type=str, required=True, help="Path to the model directory."
+    required=True,
-    )
+    help="Path to the model directory.",
-    parser.add_argument(
+)
-        "--input_json_file",
+@click.option(
-        type=str,
+    "--input_json_file",
-        required=True,
+    type=click.Path(exists=True),
-        help="Path to the input JSONL file.",
+    required=True,
-    )
+    help="Path to the input JSONL file.",
-    parser.add_argument(
+)
-        "--output_json_file",
+@click.option(
-        type=str,
+    "--output_json_file",
-        required=True,
+    type=click.Path(),
-        help="Path to the output JSONL file.",
+    required=True,
-    )
+    help="Path to the output JSONL file.",
-    parser.add_argument(
+)
-        "--question_key",
+@click.option(
-        type=str,
+    "--question_key", default="question", help="Key for the question in input JSON."
-        default="question",
+)
-        help="Key for the question in the input JSON (default: question).",
+@click.option(
-    )
+    "--response_key", default="response", help="Key for the response in output JSON."
-    parser.add_argument(
+)
-        "--response_key",
+@click.option("--temperature", type=float, default=0.60, help="Sampling temperature.")
-        type=str,
+@click.option("--top_k", type=int, default=30, help="Top-k filtering.")
-        default="response",
+@click.option("--top_p", type=float, default=0.95, help="Top-p filtering.")
-        help="Key for the response in the output JSON (default: response).",
+@click.option("--batch_size", type=int, default=1, help="Batch size.")
-    )
+@click.option("--num_samples", type=int, default=1, help="Responses per prompt.")
-    parser.add_argument(
+@click.option("--max_tokens", type=int, default=None, help="Max tokens to generate.")
-        "--temperature",
+@click.option("--cache_len", type=int, default=2048, help="KV cache length.")
-        type=float,
+@click.option("--frequency_penalty", type=float, default=0.0, help="Frequency penalty.")
-        default=0.60,
+@click.option(
-        help="Temperature for generating responses (default: 0.60).",
+    "--rep_window", type=int, default=64, help="Window size for frequency penalty."
-    )
+)
-    parser.add_argument(
+def generate_command(**kwargs):
-        "--top_k",
+    """Batch generation from a JSONL prompt file."""
-        type=int,
-        default=30,
-        help="Top-k value for generating responses (default: 30).",
-    )
-    parser.add_argument(
-        "--top_p",
-        type=float,
-        default=0.95,
-        help="Top-p value for generating responses (default: 0.95).",
-    )
-    parser.add_argument(
-        "--batch_size",
-        type=int,
-        default=1,
-        help="Batch size for generating responses (default: 1).",
-    )
-    parser.add_argument(
-        "--num_samples",
-        type=int,
-        default=1,
-        help="Number of responses per prompt (expands batch internally, default: 1).",
-    )
-    parser.add_argument(
-        "--max_tokens",
-        type=int,
-        default=None,
-        help=(
-            "Maximum tokens to generate "
-            "(default: model config max_position_embeddings)."
-        ),
-    )
-    parser.add_argument(
-        "--cache_len",
-        type=int,
-        default=2048,
-        help="KV cache & prompt truncation length (default: 2048, lower = less memory).",
-    )
-    parser.add_argument(
-        "--frequency_penalty",
-        type=float,
-        default=0.0,
-        help="Frequency penalty to reduce repetition (default: 0.0, try 0.5-1.0).",
-    )
-    parser.add_argument(
-        "--rep_window",
-        type=int,
-        default=64,
-        help="Window size for frequency penalty (default: 64).",
-    )
-
-    args = parser.parse_args()
-
     with torch.inference_mode():
-        processor(**vars(args))
+        processor(**kwargs)
+
+
+if __name__ == "__main__":
+    generate_command()

scripts/tools/preprocess.py

@@ -1,48 +1,50 @@
 """CLI: JSONL → tokenized .h5/.bin via config-driven Pipeline."""
 
-import argparse
+import click
 
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.pipeline import Pipeline
 
 
-def main():
+@click.command(
-    parser = argparse.ArgumentParser(
+    name="preprocess", help="Tokenize and pack raw JSONL data into .bin/.h5 format."
-        description="Raw JSONL → tokenized .h5/.bin via config-driven Pipeline"
+)
-    )
+@click.argument("inputs", nargs=-1, type=click.Path(exists=True), required=True)
-    parser.add_argument(
+@click.option(
-        "inputs", nargs="+", metavar="JSONL", help="One or more JSONL files"
+    "--output_dir", "-o", type=click.Path(), required=True, help="Output directory."
-    )
+)
-    parser.add_argument("--output_dir", "-o", required=True, help="Output directory")
+@click.option(
-    parser.add_argument(
+    "--config",
-        "--config", "-c", required=True, help="Path to pipeline config JSON"
+    "-c",
-    )
+    "pipeline_config",
-    parser.add_argument(
+    type=click.Path(exists=True),
-        "--tokenizer_path",
+    required=True,
-        default="params",
+    help="Pipeline config JSON.",
-        help="Path to tokenizer directory (default: params)",
+)
-    )
+@click.option(
-    parser.add_argument(
+    "--tokenizer_path",
-        "--batch_size",
+    type=click.Path(exists=True),
-        type=int,
+    default="params",
-        default=None,
+    help="Path to tokenizer directory.",
-        help="Number of records tokenized together (default: config value)",
+)
-    )
+@click.option("--batch_size", type=int, default=None, help="Records per batch.")
-    args = parser.parse_args()
+def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batch_size):
-
+    """Tokenize and pack raw JSONL data into .bin/.h5 format."""
-    config = PipelineConfig.from_file(args.config)
+    config = PipelineConfig.from_file(pipeline_config)
-    if args.batch_size is not None:
+    if batch_size is not None:
-        if args.batch_size < 1:
+        if batch_size < 1:
-            parser.error("--batch_size must be at least 1")
+            raise click.BadParameter("--batch_size must be at least 1")
-        config.preprocessing.batch_size = args.batch_size
+        config.preprocessing.batch_size = batch_size
 
+    click.echo(f"Preprocessing {len(inputs)} file(s) → {output_dir}")
     Pipeline(
         config=config,
-        input_paths=args.inputs,
+        input_paths=list(inputs),
-        output_dir=args.output_dir,
+        output_dir=output_dir,
-        tokenizer_path=args.tokenizer_path,
+        tokenizer_path=tokenizer_path,
     ).run()
+    click.echo("Done.")
 
 
 if __name__ == "__main__":
-    main()
+    preprocess_command()

scripts/tools/server.py

@@ -1,72 +1,58 @@
-import argparse
 from pathlib import Path
 
+import click
 import torch
 
 from astrai.inference import run_server
 
+_DTYPES = ["bfloat16", "float16", "float32"]
 
-def main():
-    parser = argparse.ArgumentParser(description="Start AstrAI inference HTTP server")
-    parser.add_argument(
-        "--host", default="0.0.0.0", help="Host address (default: 0.0.0.0)"
-    )
-    parser.add_argument(
-        "--port", type=int, default=8000, help="Port number (default: 8000)"
-    )
-    parser.add_argument(
-        "--reload", action="store_true", help="Enable auto-reload for development"
-    )
-    parser.add_argument(
-        "--param_path",
-        type=Path,
-        default=None,
-        help="Path to model parameters (default: project_root/params)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda",
-        help="Device to load model on (default: cuda)",
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        default="bfloat16",
-        choices=["bfloat16", "float16", "float32"],
-        help="Data type for model weights (default: bfloat16)",
-    )
-    parser.add_argument(
-        "--max_batch_size",
-        type=int,
-        default=16,
-        help="Maximum batch size for continuous batching (default: 16)",
-    )
-    args = parser.parse_args()
 
-    # Convert dtype string to torch dtype
+@click.command(name="serve", help="Launch inference server (OpenAI-compatible API).")
+@click.option("--host", default="0.0.0.0", help="Host address.")
+@click.option("--port", type=int, default=8000, help="Port number.")
+@click.option("--reload", is_flag=True, help="Enable auto-reload for development.")
+@click.option(
+    "--param_path",
+    type=click.Path(exists=True),
+    default=None,
+    help="Path to model parameters.",
+)
+@click.option("--device", default="cuda", help="Device to load model on.")
+@click.option(
+    "--dtype",
+    type=click.Choice(_DTYPES),
+    default="bfloat16",
+    help="Data type for model weights.",
+)
+@click.option(
+    "--max_batch_size",
+    type=int,
+    default=16,
+    help="Maximum batch size for continuous batching.",
+)
+def server_command(host, port, reload, param_path, device, dtype, max_batch_size):
+    """Launch inference server (OpenAI-compatible API)."""
     dtype_map = {
         "bfloat16": torch.bfloat16,
         "float16": torch.float16,
         "float32": torch.float32,
     }
-    dtype = dtype_map[args.dtype]
-
     project_root = Path(__file__).parent.parent.parent
-    param_path = args.param_path or (project_root / "params")
+    param_path = param_path or str(project_root / "params")
-    print(f"Starting AstrAI inference server on http://{args.host}:{args.port}")
+
-    print(f"Model parameters expected at: {param_path}")
+    click.echo(f"Starting server on http://{host}:{port}")
-    print(f"Device: {args.device}, Dtype: {args.dtype}")
+    click.echo(f"Model: {param_path}  |  Device: {device}  |  Dtype: {dtype}")
     run_server(
-        host=args.host,
+        host=host,
-        port=args.port,
+        port=port,
-        reload=args.reload,
+        reload=reload,
-        device=args.device,
+        device=device,
-        dtype=dtype,
+        dtype=dtype_map[dtype],
-        param_path=param_path,
+        param_path=Path(param_path),
-        max_batch_size=args.max_batch_size,
+        max_batch_size=max_batch_size,
     )
 
 
 if __name__ == "__main__":
-    main()
+    server_command()

scripts/tools/train.py

@@ -1,11 +1,11 @@
-import argparse
 import os
+from collections.abc import Callable
 from functools import partial
-from typing import Any, Callable, Dict, Optional
+from typing import Any
 
+import click
 import torch
-import torch.optim as optim
+from torch import Tensor, nn, optim
-from torch import Tensor, nn
 
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
@@ -28,14 +28,14 @@ class MuonMix(optim.Optimizer):
         ns_steps: int = 5,
         adjust_lr_fn: str = "match_rms_adamw",
     ):
-        defaults = dict(
+        defaults = {
-            lr=lr,
+            "lr": lr,
-            weight_decay=weight_decay,
+            "weight_decay": weight_decay,
-            momentum=momentum,
+            "momentum": momentum,
-            nesterov=nesterov,
+            "nesterov": nesterov,
-            ns_steps=ns_steps,
+            "ns_steps": ns_steps,
-            adjust_lr_fn=adjust_lr_fn,
+            "adjust_lr_fn": adjust_lr_fn,
-        )
+        }
         params = [p for p in model.parameters() if p.requires_grad]
         super().__init__(params, defaults)
 
@@ -82,312 +82,252 @@ class MuonMix(optim.Optimizer):
         self.muon.zero_grad(set_to_none)
         self.adamw.zero_grad(set_to_none)
 
-    def state_dict(self) -> Dict[str, Any]:
+    def state_dict(self) -> dict[str, Any]:
         return {
             "muon": self.muon.state_dict(),
             "adamw": self.adamw.state_dict(),
         }
 
-    def load_state_dict(self, state_dict: Dict[str, Any]):
+    def load_state_dict(self, state_dict: dict[str, Any]):
         self.muon.load_state_dict(state_dict["muon"])
         self.adamw.load_state_dict(state_dict["adamw"])
         self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
 
 
-def parse_args() -> argparse.Namespace:
+def _merge_yaml_into_kwargs(config_path: str, passed_kwargs: dict) -> dict:
+    """Load YAML config, then override with explicit CLI kwargs (None excluded)."""
+    import yaml
 
-    parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
+    with open(config_path) as f:
+        cfg = yaml.safe_load(f)
 
-    parser.add_argument(
+    merged = {}
-        "--train_type",
+    for section in ("model", "data", "parallel", "training", "ckpt", "log"):
-        type=str,
+        if section in cfg:
-        required=True,
+            merged.update(cfg[section])
-        choices=["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"],
-        help="Train type.",
-    )
-    parser.add_argument(
-        "--data_root_path",
-        type=str,
-        required=True,
-        help="Path to the root directory of the dataset.",
-    )
-    parser.add_argument(
-        "--param_path",
-        type=str,
-        required=True,
-        help="Path to the model parameters or resume checkpoint.",
-    )
-    parser.add_argument(
-        "--resume",
-        action="store_true",
-        default=False,
-        help="Resume training from checkpoint at --param_path "
-        "(restore epoch, consumed_samples, optimizer & scheduler state).",
-    )
 
-    parser.add_argument(
+    for key, value in passed_kwargs.items():
-        "--n_epoch", type=int, default=1, help="Number of epochs to train."
+        if value is not None:
-    )
+            merged[key] = value
-    parser.add_argument(
-        "--batch_per_device", type=int, default=1, help="Batch size per GPU."
-    )
-    parser.add_argument(
-        "--grad_accum_steps",
-        type=int,
-        default=1,
-        help="Number of iterations between each optimizer step.",
-    )
-    parser.add_argument(
-        "--warmup_ratio",
-        type=float,
-        default=0.05,
-        help="Fraction of total steps used for LR warmup.",
-    )
-    parser.add_argument(
-        "--max_lr", type=float, default=3e-4, help="Max learning rate for training."
-    )
-    parser.add_argument(
-        "--max_grad_norm",
-        type=float,
-        default=1.0,
-        help="Max gradient norm for clipping. None disables clipping.",
-    )
-    parser.add_argument(
-        "--weight_decay",
-        type=float,
-        default=0.1,
-        help="Weight decay (applied to Muon matrix params; non-matrix use 0).",
-    )
-    parser.add_argument(
-        "--muon_momentum",
-        type=float,
-        default=0.95,
-        help="Momentum factor for Muon optimizer.",
-    )
-    parser.add_argument(
-        "--muon_nesterov",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Enable Nesterov momentum for Muon.",
-    )
-    parser.add_argument(
-        "--muon_ns_steps",
-        type=int,
-        default=5,
-        help="Newton-Schulz iteration steps for Muon.",
-    )
-    parser.add_argument(
-        "--muon_adjust_lr",
-        type=str,
-        default="match_rms_adamw",
-        choices=["original", "match_rms_adamw"],
-        help="Muon learning rate adjustment strategy.",
-    )
-    parser.add_argument(
-        "--random_seed", type=int, default=3407, help="Random seed for reproducibility."
-    )
-    parser.add_argument(
-        "--num_workers", type=int, default=4, help="Number of workers for data loading."
-    )
-    parser.add_argument(
-        "--no_pin_memory",
-        action="store_false",
-        dest="pin_memory",
-        help="Disable pin memory",
-    )
-    parser.add_argument(
-        "--window_size",
-        type=int,
-        default=None,
-        help="Max length of the input sequence.",
-    )
-    parser.add_argument(
-        "--stride", type=int, default=None, help="Step size of the input sequence."
-    )
-    parser.add_argument("--dpo_beta", type=float, default=0.1, help="DPO beta value.")
-    parser.add_argument("--group_size", type=int, default=4, help="GRPO group size.")
-    parser.add_argument(
-        "--grpo_clip_eps", type=float, default=0.2, help="GRPO clipping epsilon."
-    )
-    parser.add_argument(
-        "--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient."
-    )
-    parser.add_argument(
-        "--label_smoothing",
-        type=float,
-        default=0.0,
-        help="cross_entropy function label smoothing parameter",
-    )
 
-    # online rollout
+    return merged
-    parser.add_argument(
-        "--rollout_interval",
-        type=int,
-        default=512,
-        help="Number of optimizer steps between online rollouts.",
-    )
-    parser.add_argument(
-        "--rollout_temperature",
-        type=float,
-        default=0.7,
-        help="Sampling temperature for online rollout.",
-    )
-    parser.add_argument(
-        "--rollout_top_k",
-        type=int,
-        default=0,
-        help="Top-k filtering for online rollout (0=disable).",
-    )
-    parser.add_argument(
-        "--rollout_top_p",
-        type=float,
-        default=0.9,
-        help="Top-p (nucleus) filtering for online rollout.",
-    )
-    parser.add_argument(
-        "--rollout_max_tokens",
-        type=int,
-        default=1024,
-        help="Maximum generated tokens per response in rollout.",
-    )
 
-    parser.add_argument(
-        "--gradient_checkpointing",
-        action=argparse.BooleanOptionalAction,
-        default=False,
-        help="Enable activation checkpointing for DecoderBlock modules.",
-    )
 
-    parser.add_argument(
+_TRAIN_TYPE = ["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"]
-        "--ckpt_interval",
+_PARALLEL = ["none", "ddp", "fsdp", "fsdp2"]
-        type=int,
+_SCHEDULES = ["cosine", "sgdr", "wsd"]
-        default=5000,
+_BACKENDS = ["nccl", "gloo"]
-        help="Number of iters between checkpoints.",
+_START_METHODS = ["spawn", "fork", "forkserver"]
-    )
-    parser.add_argument(
-        "--ckpt_dir",
-        type=str,
-        default="checkpoint",
-        help="Directory to save checkpoints.",
-    )
-    parser.add_argument(
-        "--val_split",
-        type=float,
-        default=None,
-        help="Ratio to split from training dataset for validation (e.g. 0.05).",
-    )
-    parser.add_argument(
-        "--val_step",
-        type=int,
-        default=1000,
-        help="Number of optimizer steps between validation runs.",
-    )
-    parser.add_argument(
-        "--metrics",
-        nargs="*",
-        default=["loss", "lr", "grad_norm"],
-        help="Metrics to log (e.g. --metrics loss lr val_loss). Default: loss lr grad_norm.",
-    )
-    parser.add_argument(
-        "--log_dir",
-        type=str,
-        default="checkpoint/logs",
-        help="Directory for metric logs.",
-    )
-    parser.add_argument(
-        "--start_epoch", type=int, default=0, help="Start epoch for training."
-    )
-    parser.add_argument(
-        "--start_samples",
-        type=int,
-        default=0,
-        help="Start samples (per rank) for training.",
-    )
 
-    parser.add_argument(
-        "--master_addr",
-        type=str,
-        default="localhost",
-        help="Master node address for distributed training.",
-    )
-    parser.add_argument(
-        "--master_port",
-        type=str,
-        default="29500",
-        help="Master node port for distributed training.",
-    )
-    parser.add_argument(
-        "--backend",
-        type=str,
-        default="nccl",
-        help="Distributed training backend.",
-    )
-    parser.add_argument("--nprocs", type=int, default=1, help="Number of GPUs to use.")
-    parser.add_argument(
-        "--parallel_mode",
-        type=str,
-        default="none",
-        choices=["none", "ddp", "fsdp", "fsdp2"],
-        help="Parallel training strategy (none, ddp, fsdp, fsdp2).",
-    )
-    parser.add_argument(
-        "--device_type", type=str, default="cuda", help="Device type to use."
-    )
-    parser.add_argument(
-        "--start_method",
-        type=str,
-        default="spawn",
-        choices=["spawn", "fork", "forkserver"],
-        help="Multiprocessing start method.",
-    )
-    parser.add_argument(
-        "--neftune_alpha",
-        type=float,
-        default=0.0,
-        help="NEFTune noise alpha (0=disabled, typical: 5.0).",
-    )
 
-    parser.add_argument(
+@click.command(
-        "--schedule_type",
+    name="train",
-        type=str,
+    help="Start model training (pretrain / SFT / DPO / GRPO).",
-        default="cosine",
+    context_settings={"show_default": True},
-        choices=["cosine", "sgdr", "wsd"],
+)
-        help="Learning rate scheduler type.",
+@click.option(
-    )
+    "--config",
-    parser.add_argument(
+    "-c",
-        "--min_rate",
+    "config_path",
-        type=float,
+    type=click.Path(exists=True),
-        default=None,
+    help="YAML config file. CLI flags override YAML values.",
-        help="Minimum LR as fraction of base LR. Uses scheduler default if not set (cosine/sgdr: 0.05, wsd: 0.0).",
+)
-    )
+@click.option(
-    parser.add_argument(
+    "--train_type",
-        "--cycle_length",
+    type=click.Choice(_TRAIN_TYPE),
-        type=int,
+    required=False,
-        default=None,
+    help="Training type.",
-        help="SGDR first cycle length in steps. Defaults to total_steps - warmup_steps.",
+)
-    )
+@click.option(
-    parser.add_argument(
+    "--data_root_path",
-        "--t_mult",
+    type=click.Path(exists=True),
-        type=int,
+    help="Root directory of the dataset.",
-        default=2,
+)
-        help="SGDR cycle length multiplier per restart.",
+@click.option(
-    )
+    "--param_path",
-    parser.add_argument(
+    type=click.Path(exists=True),
-        "--stable_steps",
+    help="Path to model parameters or resume checkpoint.",
-        type=int,
+)
-        default=None,
+@click.option("--resume", is_flag=True, default=False, help="Resume from checkpoint.")
-        help="WSD stable plateau steps. Required when --schedule_type wsd.",
+@click.option("--n_epoch", type=int, default=1, help="Number of epochs.")
-    )
+@click.option("--batch_per_device", type=int, default=1, help="Batch size per GPU.")
-    parser.add_argument(
+@click.option(
-        "--decay_steps",
+    "--grad_accum_steps", type=int, default=1, help="Gradient accumulation steps."
-        type=int,
+)
-        default=None,
+@click.option(
-        help="WSD decay steps. Defaults to total_steps - warmup_steps - stable_steps.",
+    "--warmup_ratio",
-    )
+    type=float,
+    default=0.05,
+    help="Fraction of total steps for LR warmup.",
+)
+@click.option("--max_lr", type=float, default=3e-4, help="Max learning rate.")
+@click.option(
+    "--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping."
+)
+@click.option("--weight_decay", type=float, default=0.1, help="Weight decay.")
+@click.option("--muon_momentum", type=float, default=0.95, help="Muon momentum factor.")
+@click.option("--muon_nesterov/--no-muon_nesterov", default=True, help="Muon Nesterov.")
+@click.option("--muon_ns_steps", type=int, default=5, help="Muon Newton-Schulz steps.")
+@click.option(
+    "--muon_adjust_lr",
+    type=click.Choice(["original", "match_rms_adamw"]),
+    default="match_rms_adamw",
+    help="Muon LR adjustment strategy.",
+)
+@click.option("--random_seed", type=int, default=3407, help="Random seed.")
+@click.option("--num_workers", type=int, default=4, help="DataLoader workers.")
+@click.option("--pin_memory/--no-pin_memory", default=True, help="Pin memory.")
+@click.option(
+    "--window_size", type=int, default=None, help="Max input sequence length."
+)
+@click.option("--stride", type=int, default=None, help="Step size for sliding window.")
+@click.option("--dpo_beta", type=float, default=0.1, help="DPO beta.")
+@click.option("--group_size", type=int, default=4, help="GRPO group size.")
+@click.option("--grpo_clip_eps", type=float, default=0.2, help="GRPO clip epsilon.")
+@click.option(
+    "--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient."
+)
+@click.option("--label_smoothing", type=float, default=0.0, help="Label smoothing.")
+@click.option(
+    "--rollout_interval", type=int, default=512, help="Steps between rollouts."
+)
+@click.option(
+    "--rollout_temperature", type=float, default=0.7, help="Rollout temperature."
+)
+@click.option("--rollout_top_k", type=int, default=0, help="Rollout top-k (0=disable).")
+@click.option("--rollout_top_p", type=float, default=0.9, help="Rollout top-p.")
+@click.option(
+    "--rollout_max_tokens",
+    type=int,
+    default=1024,
+    help="Max tokens per rollout response.",
+)
+@click.option(
+    "--gradient_checkpointing/--no-gradient_checkpointing",
+    default=False,
+    help="Enable activation checkpointing.",
+)
+@click.option(
+    "--ckpt_interval", type=int, default=5000, help="Steps between checkpoints."
+)
+@click.option(
+    "--ckpt_dir", type=click.Path(), default="checkpoint", help="Checkpoint directory."
+)
+@click.option("--val_split", type=float, default=None, help="Validation split ratio.")
+@click.option(
+    "--val_step", type=int, default=1000, help="Steps between validation runs."
+)
+@click.option(
+    "--metrics",
+    multiple=True,
+    default=("loss", "lr", "grad_norm"),
+    help="Metrics to log (repeatable).",
+)
+@click.option(
+    "--log_dir",
+    type=click.Path(),
+    default="checkpoint/logs",
+    help="Directory for metric logs.",
+)
+@click.option("--start_epoch", type=int, default=0, help="Start epoch.")
+@click.option("--start_samples", type=int, default=0, help="Start samples (per rank).")
+@click.option(
+    "--master_addr", type=str, default="localhost", help="Master node address."
+)
+@click.option("--master_port", type=str, default="29500", help="Master node port.")
+@click.option(
+    "--backend",
+    type=click.Choice(_BACKENDS),
+    default="nccl",
+    help="Distributed backend.",
+)
+@click.option("--nprocs", type=int, default=1, help="Number of GPUs.")
+@click.option(
+    "--parallel_mode",
+    type=click.Choice(_PARALLEL),
+    default="none",
+    help="Parallel strategy.",
+)
+@click.option("--device_type", type=str, default="cuda", help="Device type.")
+@click.option(
+    "--start_method",
+    type=click.Choice(_START_METHODS),
+    default="spawn",
+    help="Multiprocessing start method.",
+)
+@click.option("--neftune_alpha", type=float, default=0.0, help="NEFTune noise alpha.")
+@click.option(
+    "--schedule_type",
+    type=click.Choice(_SCHEDULES),
+    default="cosine",
+    help="LR scheduler.",
+)
+@click.option(
+    "--min_rate", type=float, default=None, help="Minimum LR as fraction of base LR."
+)
+@click.option("--cycle_length", type=int, default=None, help="SGDR first cycle length.")
+@click.option("--t_mult", type=int, default=2, help="SGDR cycle length multiplier.")
+@click.option(
+    "--stable_steps", type=int, default=None, help="WSD stable plateau steps."
+)
+@click.option("--decay_steps", type=int, default=None, help="WSD decay steps.")
+@click.option("--tp_size", type=int, default=None, help="Tensor parallelism (future).")
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    default=False,
+    help="Validate config and print plan, do not train.",
+)
+@click.pass_context
+def train_command(ctx, config_path, dry_run, metrics, **kwargs):
+    """Start model training (pretrain / SFT / DPO / GRPO)."""
+    if config_path:
+        kwargs = _merge_yaml_into_kwargs(config_path, kwargs)
 
-    args = parser.parse_args()
+    required = ["train_type", "data_root_path", "param_path"]
+    missing = [k for k in required if kwargs.get(k) is None]
+    if missing:
+        raise click.UsageError(
+            f"Missing required options: {', '.join(missing)}. "
+            f"Use --config YAML or provide them directly."
+        )
 
-    return args
+    # Convert tuple back to list
+    kwargs["metrics"] = list(metrics)
+    # Remove tp_size (not yet wired)
+    kwargs.pop("tp_size", None)
+
+    if dry_run:
+        _print_dry_run(kwargs)
+        return
+
+    train(**kwargs)
+
+
+def _print_dry_run(kwargs: dict) -> None:
+    """Print training plan summary."""
+    rows = [
+        ("Train type", kwargs.get("train_type")),
+        ("Model path", kwargs.get("param_path")),
+        ("Data path", kwargs.get("data_root_path")),
+        ("Parallel mode", kwargs.get("parallel_mode", "none")),
+        ("GPUs", str(kwargs.get("nprocs", 1))),
+        ("Epochs", str(kwargs.get("n_epoch", 1))),
+        ("Batch/device", str(kwargs.get("batch_per_device", 1))),
+        ("Grad accum", str(kwargs.get("grad_accum_steps", 1))),
+        ("Max LR", str(kwargs.get("max_lr", "?"))),
+        ("Schedule", str(kwargs.get("schedule_type", "cosine"))),
+        ("Warmup ratio", str(kwargs.get("warmup_ratio", 0.05))),
+        ("Window size", str(kwargs.get("window_size", "config default"))),
+        ("Checkpoint dir", str(kwargs.get("ckpt_dir", "checkpoint"))),
+        ("Checkpoint interval", str(kwargs.get("ckpt_interval", 5000))),
+        ("Resume", str(kwargs.get("resume", False))),
+    ]
+    max_len = max(len(k) for k, _ in rows)
+    click.secho("\n=== Training Plan (dry-run) ===", fg="cyan", bold=True)
+    for key, val in rows:
+        click.echo(f"  {key:<{max_len}s} : {val}")
+    click.secho("=" * 40, fg="cyan")
 
 
 def create_model(config):
@@ -497,7 +437,7 @@ def train(
     rollout_top_k = kwargs.pop("rollout_top_k", 0)
     rollout_top_p = kwargs.pop("rollout_top_p", 0.9)
     rollout_max_tokens = kwargs.pop("rollout_max_tokens", 1024)
-    reward_model_fn: Optional[Callable[[], BaseRewardModel]] = None
+    reward_model_fn: Callable[[], BaseRewardModel] | None = None
 
     executor_kwargs = {}
     if parallel_mode == "ddp":
@@ -611,5 +551,4 @@ def train(
 
 
 if __name__ == "__main__":
-    args = parse_args()
+    train_command()
-    train(**vars(args))