tinymind-native-colab-handoff/bundle/train/cli_light.py

"""Lightweight CLI parser helpers for tests and dependency-sensitive checks.

The full ``train.cli`` module intentionally wires many dataset/evaluation
commands, some of which import native extensions such as pyarrow. Unit tests
that only need parser registration should use this module to avoid loading the
full training/evaluation stack.
"""

from __future__ import annotations

import argparse


def build_axiomflow_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="tinymind-light")
    sub = parser.add_subparsers(dest="command", required=True)
    axiomflow = sub.add_parser(
        "axiomflow-bench",
        help="Run AxiomFlow HyperWeave bounded-memory reference benchmark",
    )
    axiomflow.add_argument("--out-dir", default="reports/axiomflow")
    axiomflow.add_argument("--dim", type=int, default=128)
    axiomflow.add_argument("--seq-len", type=int, default=64)
    axiomflow.add_argument("--batch-size", type=int, default=1)
    axiomflow.add_argument("--local-window", type=int, default=16)
    axiomflow.add_argument("--memory-slots", type=int, default=8)
    axiomflow.add_argument("--memory-rank", type=int, default=16)
    axiomflow.add_argument("--retrieval-top-k", type=int, default=4)
    axiomflow.add_argument("--retrieved-chunk-tokens", type=int, default=32)
    axiomflow.add_argument("--seed", type=int, default=20260525)
    return parser