autogen typer options from pydantic models
This commit is contained in:
Dan Saunders
parent
356fb08b03
commit
cf966fe98e
2 changed files with 151 additions and 109 deletions
117
cli.py
117
cli.py
|
|
@ -2,14 +2,12 @@ import logging
|
|||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional, List
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
|
||||
from cli.config import Config, load_config
|
||||
|
||||
# CLI args that should not be passed to cfg.apply_overrides()
|
||||
_EXCLUDED_CLI_ARGS = ("config", "dry_run", "verbose", "hf_token", "cfg")
|
||||
from cli.options import add_options_from_config
|
||||
|
||||
app = typer.Typer(
|
||||
help="Command-line interface for Unsloth training, chat, and export.",
|
||||
|
|
@ -27,122 +25,24 @@ def configure_logging(verbose: bool):
|
|||
|
||||
|
||||
@app.command()
|
||||
@add_options_from_config(Config)
|
||||
def train(
|
||||
model: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--model",
|
||||
"-m",
|
||||
help="HF model id or local path. Required unless provided in --config.",
|
||||
),
|
||||
training_type: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--training-type",
|
||||
help="Training mode: 'lora' (LoRA/QLoRA) or 'full'. Defaults to 'lora'.",
|
||||
),
|
||||
hf_token: Optional[str] = typer.Option(
|
||||
None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed."
|
||||
),
|
||||
max_seq_length: Optional[int] = typer.Option(None, "--max-seq-length"),
|
||||
load_in_4bit: Optional[bool] = typer.Option(
|
||||
None, "--load-in-4bit/--no-load-in-4bit"
|
||||
),
|
||||
output_dir: Optional[Path] = typer.Option(
|
||||
None,
|
||||
"--output-dir",
|
||||
help="Where to store checkpoints. Defaults to ./outputs",
|
||||
),
|
||||
dataset: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--dataset",
|
||||
"-d",
|
||||
help="HF dataset to train on (e.g. 'tatsu-lab/alpaca').",
|
||||
),
|
||||
local_dataset: Optional[List[str]] = typer.Option(
|
||||
None,
|
||||
"--local-dataset",
|
||||
help="Filename(s) under datasets/ to use (e.g. 'alpaca_unsloth.json').",
|
||||
),
|
||||
format_type: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--format-type",
|
||||
help="Dataset formatting: auto|alpaca|chatml|sharegpt. Defaults to auto.",
|
||||
),
|
||||
num_epochs: Optional[int] = typer.Option(None, "--epochs"),
|
||||
learning_rate: Optional[float] = typer.Option(None, "--lr"),
|
||||
batch_size: Optional[int] = typer.Option(None, "--batch-size"),
|
||||
gradient_accumulation_steps: Optional[int] = typer.Option(None, "--grad-accum"),
|
||||
warmup_steps: Optional[int] = typer.Option(None, "--warmup-steps"),
|
||||
max_steps: Optional[int] = typer.Option(
|
||||
None, "--max-steps", help="Overrides epochs if >0."
|
||||
),
|
||||
save_steps: Optional[int] = typer.Option(
|
||||
None, "--save-steps", help="0 uses trainer defaults."
|
||||
),
|
||||
weight_decay: Optional[float] = typer.Option(None, "--weight-decay"),
|
||||
random_seed: Optional[int] = typer.Option(None, "--seed"),
|
||||
packing: Optional[bool] = typer.Option(None, "--packing/--no-packing"),
|
||||
train_on_completions: Optional[bool] = typer.Option(
|
||||
None, "--train-on-completions", help="Train on responses only when supported."
|
||||
),
|
||||
lora_r: Optional[int] = typer.Option(None, "--lora-r"),
|
||||
lora_alpha: Optional[int] = typer.Option(None, "--lora-alpha"),
|
||||
lora_dropout: Optional[float] = typer.Option(None, "--lora-dropout"),
|
||||
gradient_checkpointing: Optional[bool] = typer.Option(
|
||||
None, "--gradient-checkpointing/--no-gradient-checkpointing"
|
||||
),
|
||||
target_modules: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--target-modules",
|
||||
help="Comma-separated target modules for LoRA.",
|
||||
),
|
||||
vision_all_linear: Optional[bool] = typer.Option(
|
||||
None,
|
||||
"--vision-all-linear/--no-vision-all-linear",
|
||||
help="For vision models, finetune all linear layers (mirrors UI toggle).",
|
||||
),
|
||||
finetune_vision_layers: Optional[bool] = typer.Option(
|
||||
None,
|
||||
"--finetune-vision-layers/--no-finetune-vision-layers",
|
||||
help="For vision LoRA: train vision layers.",
|
||||
),
|
||||
finetune_language_layers: Optional[bool] = typer.Option(
|
||||
None,
|
||||
"--finetune-language-layers/--no-finetune-language-layers",
|
||||
help="For vision LoRA: train language layers.",
|
||||
),
|
||||
finetune_attention_modules: Optional[bool] = typer.Option(
|
||||
None,
|
||||
"--finetune-attention-modules/--no-finetune-attention-modules",
|
||||
help="For vision LoRA: train attention modules.",
|
||||
),
|
||||
finetune_mlp_modules: Optional[bool] = typer.Option(
|
||||
None,
|
||||
"--finetune-mlp-modules/--no-finetune-mlp-modules",
|
||||
help="For vision LoRA: train MLP modules.",
|
||||
),
|
||||
use_rslora: Optional[bool] = typer.Option(None, "--rslora/--no-rslora"),
|
||||
use_loftq: Optional[bool] = typer.Option(None, "--loftq/--no-loftq"),
|
||||
enable_wandb: Optional[bool] = typer.Option(None, "--wandb/--no-wandb"),
|
||||
wandb_project: Optional[str] = typer.Option(None, "--wandb-project"),
|
||||
wandb_token: Optional[str] = typer.Option(
|
||||
None, "--wandb-token", envvar="WANDB_API_KEY"
|
||||
),
|
||||
enable_tensorboard: Optional[bool] = typer.Option(
|
||||
None, "--tensorboard/--no-tensorboard", help="Enable TensorBoard logging."
|
||||
),
|
||||
tensorboard_dir: Optional[str] = typer.Option(None, "--tensorboard-dir"),
|
||||
config: Optional[Path] = typer.Option(
|
||||
None,
|
||||
"--config",
|
||||
"-c",
|
||||
help="Path to YAML/JSON config file. CLI flags override config values.",
|
||||
),
|
||||
hf_token: Optional[str] = typer.Option(
|
||||
None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed."
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
False,
|
||||
"--dry-run",
|
||||
help="Show resolved config and exit without training.",
|
||||
),
|
||||
verbose: bool = typer.Option(False, "--verbose/--quiet"),
|
||||
config_overrides: dict = None, # Injected by decorator
|
||||
):
|
||||
"""
|
||||
Launch training using the existing Unsloth training backend.
|
||||
|
|
@ -154,8 +54,7 @@ def train(
|
|||
raise typer.Exit(code=2)
|
||||
|
||||
# Apply CLI overrides
|
||||
cli_args = {k: v for k, v in locals().items() if k not in _EXCLUDED_CLI_ARGS}
|
||||
cfg.apply_overrides(**cli_args)
|
||||
cfg.apply_overrides(**config_overrides)
|
||||
|
||||
# Dry run: show resolved config and exit
|
||||
if dry_run:
|
||||
|
|
|
|||
143
cli/options.py
Normal file
143
cli/options.py
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
"""Generate Typer CLI options from Pydantic models."""
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional, get_args, get_origin
|
||||
|
||||
import typer
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
def _python_name_to_cli_flag(name: str) -> str:
|
||||
"""Convert python_name to --cli-flag."""
|
||||
return "--" + name.replace("_", "-")
|
||||
|
||||
|
||||
def _unwrap_optional(annotation: Any) -> Any:
|
||||
"""Unwrap Optional[X] to X."""
|
||||
origin = get_origin(annotation)
|
||||
if origin is not None:
|
||||
args = get_args(annotation)
|
||||
if type(None) in args:
|
||||
non_none = [a for a in args if a is not type(None)]
|
||||
if non_none:
|
||||
return non_none[0]
|
||||
return annotation
|
||||
|
||||
|
||||
def _is_bool_field(annotation: Any) -> bool:
|
||||
"""Check if field is a boolean (including Optional[bool])."""
|
||||
return _unwrap_optional(annotation) is bool
|
||||
|
||||
|
||||
def _is_list_type(annotation: Any) -> bool:
|
||||
"""Check if type is a List."""
|
||||
return get_origin(annotation) is list
|
||||
|
||||
|
||||
def _get_python_type(annotation: Any) -> type:
|
||||
"""Get the Python type for annotation."""
|
||||
unwrapped = _unwrap_optional(annotation)
|
||||
if unwrapped in (str, int, float, bool, Path):
|
||||
return unwrapped
|
||||
return str
|
||||
|
||||
|
||||
def _collect_config_fields(config_class: type[BaseModel]) -> list[tuple[str, Any]]:
|
||||
"""
|
||||
Collect all fields from a config class, flattening nested models.
|
||||
Returns list of (name, field_info) tuples.
|
||||
Raises ValueError on duplicate field names.
|
||||
"""
|
||||
fields = []
|
||||
seen_names: set[str] = set()
|
||||
|
||||
for name, field_info in config_class.model_fields.items():
|
||||
annotation = field_info.annotation
|
||||
# Skip nested models - recurse into them
|
||||
if isinstance(annotation, type) and issubclass(annotation, BaseModel):
|
||||
for nested_name, nested_field in annotation.model_fields.items():
|
||||
if nested_name in seen_names:
|
||||
raise ValueError(f"Duplicate field name '{nested_name}' in config")
|
||||
seen_names.add(nested_name)
|
||||
fields.append((nested_name, nested_field))
|
||||
else:
|
||||
if name in seen_names:
|
||||
raise ValueError(f"Duplicate field name '{name}' in config")
|
||||
seen_names.add(name)
|
||||
fields.append((name, field_info))
|
||||
return fields
|
||||
|
||||
|
||||
def add_options_from_config(config_class: type[BaseModel]) -> Callable:
|
||||
"""
|
||||
Decorator that adds CLI options for all fields in a Pydantic config model.
|
||||
|
||||
The decorated function should declare a `config_overrides: dict = None` parameter
|
||||
which will receive a dict of all CLI-provided config values.
|
||||
"""
|
||||
fields = _collect_config_fields(config_class)
|
||||
field_names = {name for name, field_info in fields if not _is_list_type(field_info.annotation)}
|
||||
|
||||
def decorator(func: Callable) -> Callable:
|
||||
sig = inspect.signature(func)
|
||||
original_params = list(sig.parameters.values())
|
||||
|
||||
# Build new parameters: config fields first, then original params
|
||||
new_params = []
|
||||
|
||||
for field_name, field_info in fields:
|
||||
annotation = field_info.annotation
|
||||
if _is_list_type(annotation):
|
||||
continue
|
||||
|
||||
flag_name = _python_name_to_cli_flag(field_name)
|
||||
help_text = field_info.description or ""
|
||||
|
||||
if _is_bool_field(annotation):
|
||||
default = typer.Option(
|
||||
None,
|
||||
f"{flag_name}/--no-{field_name.replace('_', '-')}",
|
||||
help=help_text,
|
||||
)
|
||||
param = inspect.Parameter(
|
||||
field_name,
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
default=default,
|
||||
annotation=Optional[bool],
|
||||
)
|
||||
else:
|
||||
py_type = _get_python_type(annotation)
|
||||
default = typer.Option(None, flag_name, help=help_text)
|
||||
param = inspect.Parameter(
|
||||
field_name,
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
default=default,
|
||||
annotation=Optional[py_type],
|
||||
)
|
||||
new_params.append(param)
|
||||
|
||||
# Add original params, excluding config_overrides (will be injected)
|
||||
for param in original_params:
|
||||
if param.name != "config_overrides":
|
||||
new_params.append(param)
|
||||
|
||||
new_sig = sig.replace(parameters=new_params)
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
config_overrides = {}
|
||||
for key in list(kwargs.keys()):
|
||||
if key in field_names:
|
||||
if kwargs[key] is not None:
|
||||
config_overrides[key] = kwargs[key]
|
||||
del kwargs[key]
|
||||
|
||||
kwargs["config_overrides"] = config_overrides
|
||||
return func(*args, **kwargs)
|
||||
|
||||
wrapper.__signature__ = new_sig
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
Loading…
Add table
Add a link
Reference in a new issue