diff --git a/cli.py b/cli.py index 0772a4942e..c820a8d73b 100644 --- a/cli.py +++ b/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: diff --git a/cli/options.py b/cli/options.py new file mode 100644 index 0000000000..18336efca7 --- /dev/null +++ b/cli/options.py @@ -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