Merge cli from ui-early-access and fix imports

cli.py

@@ -0,0 +1,4 @@
+from cli import app
+
+if __name__ == "__main__":
+    app()

cli/__init__.py

@@ -0,0 +1,19 @@
+import typer
+
+from cli.commands.train import train
+from cli.commands.inference import inference
+from cli.commands.export import export, list_checkpoints
+from cli.commands.ui import ui
+from cli.commands.studio import studio
+
+app = typer.Typer(
+    help="Command-line interface for Unsloth training, inference, and export.",
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+
+app.command()(train)
+app.command()(inference)
+app.command()(export)
+app.command("list-checkpoints")(list_checkpoints)
+app.command()(ui)
+app.command()(studio)

cli/commands/export.py

@@ -0,0 +1,122 @@
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+
+EXPORT_FORMATS = ["merged-16bit", "merged-4bit", "gguf", "lora"]
+GGUF_QUANTS = ["q4_k_m", "q5_k_m", "q8_0", "f16"]
+
+
+def list_checkpoints(
+    outputs_dir: Path = typer.Option(
+        Path("./outputs"), "--outputs-dir", help="Directory that holds training runs."
+    ),
+):
+    """List checkpoints detected in the outputs directory."""
+    from backend.export import ExportBackend
+
+    backend = ExportBackend()
+    checkpoints = backend.scan_checkpoints(outputs_dir=str(outputs_dir))
+    if not checkpoints:
+        typer.echo("No checkpoints found.")
+        raise typer.Exit()
+
+    for display, path in checkpoints:
+        typer.echo(f"{display}: {path}")
+
+
+def export(
+    checkpoint: Path = typer.Argument(..., help="Path to checkpoint directory."),
+    output_dir: Path = typer.Argument(..., help="Directory to save exported model."),
+    format: str = typer.Option(
+        "merged-16bit",
+        "--format",
+        "-f",
+        help=f"Export format: {', '.join(EXPORT_FORMATS)}",
+    ),
+    quantization: str = typer.Option(
+        "q4_k_m",
+        "--quantization",
+        "-q",
+        help=f"GGUF quantization method: {', '.join(GGUF_QUANTS)}",
+    ),
+    push_to_hub: bool = typer.Option(
+        False, "--push-to-hub", help="Push exported model to HuggingFace Hub."
+    ),
+    repo_id: Optional[str] = typer.Option(
+        None, "--repo-id", help="HuggingFace repo ID (username/model-name)."
+    ),
+    hf_token: Optional[str] = typer.Option(
+        None, "--hf-token", envvar="HF_TOKEN", help="HuggingFace token."
+    ),
+    private: bool = typer.Option(
+        False, "--private", help="Make the HuggingFace repo private."
+    ),
+    max_seq_length: int = typer.Option(2048, "--max-seq-length"),
+    load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
+):
+    """Export a checkpoint to various formats (merged, GGUF, LoRA adapter)."""
+    if format not in EXPORT_FORMATS:
+        typer.echo(f"Error: Invalid format '{format}'. Choose from: {', '.join(EXPORT_FORMATS)}", err=True)
+        raise typer.Exit(code=2)
+
+    if push_to_hub and not repo_id:
+        typer.echo("Error: --repo-id required when using --push-to-hub", err=True)
+        raise typer.Exit(code=2)
+
+    from backend.export import ExportBackend
+    backend = ExportBackend()
+
+    typer.echo(f"Loading checkpoint: {checkpoint}")
+    success, message = backend.load_checkpoint(
+        checkpoint_path=str(checkpoint),
+        max_seq_length=max_seq_length,
+        load_in_4bit=load_in_4bit,
+    )
+    if not success:
+        typer.echo(f"Error: {message}", err=True)
+        raise typer.Exit(code=1)
+    typer.echo(message)
+
+    typer.echo(f"Exporting as {format}...")
+    if format == "merged-16bit":
+        success, message = backend.export_merged_model(
+            save_directory=str(output_dir),
+            format_type="16-bit (FP16)",
+            push_to_hub=push_to_hub,
+            repo_id=repo_id,
+            hf_token=hf_token,
+            private=private,
+        )
+    elif format == "merged-4bit":
+        success, message = backend.export_merged_model(
+            save_directory=str(output_dir),
+            format_type="4-bit (FP4)",
+            push_to_hub=push_to_hub,
+            repo_id=repo_id,
+            hf_token=hf_token,
+            private=private,
+        )
+    elif format == "gguf":
+        success, message = backend.export_gguf(
+            save_directory=str(output_dir),
+            quantization_method=quantization.upper(),
+            push_to_hub=push_to_hub,
+            repo_id=repo_id,
+            hf_token=hf_token,
+        )
+    elif format == "lora":
+        success, message = backend.export_lora_adapter(
+            save_directory=str(output_dir),
+            push_to_hub=push_to_hub,
+            repo_id=repo_id,
+            hf_token=hf_token,
+            private=private,
+        )
+
+    if not success:
+        typer.echo(f"Error: {message}", err=True)
+        raise typer.Exit(code=1)
+
+    typer.echo(message)

cli/commands/inference.py

@@ -0,0 +1,67 @@
+import sys
+from typing import Optional
+
+import typer
+
+
+def inference(
+    model: str = typer.Argument(..., help="HF model id or local path."),
+    prompt: str = typer.Argument(..., help="Prompt to send to the model."),
+    hf_token: Optional[str] = typer.Option(
+        None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed."
+    ),
+    temperature: float = typer.Option(0.7, "--temperature"),
+    top_p: float = typer.Option(0.9, "--top-p"),
+    top_k: int = typer.Option(40, "--top-k"),
+    max_new_tokens: int = typer.Option(256, "--max-new-tokens"),
+    repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"),
+    system_prompt: str = typer.Option(
+        "",
+        "--system-prompt",
+        help="Optional system prompt to prepend.",
+    ),
+    max_seq_length: int = typer.Option(2048, "--max-seq-length"),
+    load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
+):
+    """Run a single inference using the specified model."""
+    from backend.model_config import ModelConfig
+    from backend.inference import get_inference_backend
+
+    inference_backend = get_inference_backend()
+    model_config = ModelConfig.from_ui_selection(
+        dropdown_value=model, search_value=None, hf_token=hf_token, is_lora=False
+    )
+    if not model_config:
+        typer.echo("Could not resolve model config", err=True)
+        raise typer.Exit(code=1)
+
+    if not inference_backend.load_model(
+        config=model_config,
+        max_seq_length=max_seq_length,
+        load_in_4bit=load_in_4bit,
+        hf_token=hf_token,
+    ):
+        typer.echo("Model load failed", err=True)
+        raise typer.Exit(code=1)
+
+    messages = [{"role": "user", "content": prompt}]
+    stream = inference_backend.generate_chat_response(
+        messages=messages,
+        system_prompt=system_prompt,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        max_new_tokens=max_new_tokens,
+        repetition_penalty=repetition_penalty,
+    )
+
+    typer.echo("Assistant:", nl=True)
+    previous = ""
+    for chunk in stream:
+        delta = chunk[len(previous):]
+        if delta:
+            sys.stdout.write(delta)
+            sys.stdout.flush()
+        previous = chunk
+    sys.stdout.write("\n")
+    sys.stdout.flush()

cli/commands/studio.py

@@ -0,0 +1,19 @@
+import typer
+
+
+def studio(
+    port: int = typer.Option(8000, "--port", "-p", help="Port to run the UI server on."),
+    host: str = typer.Option("0.0.0.0", "--host", "-H", help="Host address to bind to."),
+    share: bool = typer.Option(True, "--share", "-s", help="Create a public Gradio share link."),
+):
+    """Launch the Unsloth web UI for training, inference, and export."""
+    from app import demo, script_dir
+
+    typer.echo(f"Starting Unsloth UI on http://{host}:{port}")
+
+    demo.launch(
+        share=share,
+        server_port=port,
+        server_name=host,
+        favicon_path=f"{script_dir}/assets/favicon-32x32.png",
+    )

cli/commands/train.py

@@ -0,0 +1,131 @@
+import time
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+from cli.config import Config, load_config
+from cli.options import add_options_from_config
+
+
+@add_options_from_config(Config)
+def train(
+    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."
+    ),
+    wandb_token: Optional[str] = typer.Option(
+        None, "--wandb-token", envvar="WANDB_API_KEY", help="Weights & Biases API key."
+    ),
+    dry_run: bool = typer.Option(
+        False,
+        "--dry-run",
+        help="Show resolved config and exit without training.",
+    ),
+    config_overrides: dict = None,
+):
+    """Launch training using the existing Unsloth training backend."""
+    try:
+        cfg = load_config(config)
+    except FileNotFoundError as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(code=2)
+
+    cfg.apply_overrides(**config_overrides)
+
+    # CLI/env tokens take precedence over config
+    # Handle case where typer.Option isn't resolved (decorator interaction)
+    from typer.models import OptionInfo
+    if isinstance(hf_token, OptionInfo):
+        hf_token = None
+    if isinstance(wandb_token, OptionInfo):
+        wandb_token = None
+    hf_token = hf_token or cfg.logging.hf_token
+    wandb_token = wandb_token or cfg.logging.wandb_token
+
+    if dry_run:
+        import yaml
+        data = cfg.model_dump()
+        data["training"]["output_dir"] = str(data["training"]["output_dir"])
+        typer.echo(yaml.dump(data, default_flow_style=False, sort_keys=False))
+        raise typer.Exit(code=0)
+
+    if not cfg.model:
+        typer.echo("Error: provide --model or set model in --config", err=True)
+        raise typer.Exit(code=2)
+
+    if not cfg.data.dataset and not cfg.data.local_dataset:
+        typer.echo(
+            "Error: provide --dataset or --local-dataset (or via --config)", err=True
+        )
+        raise typer.Exit(code=2)
+
+    # Check if the model path is a LoRA adapter (has adapter_config.json)
+    model_path = Path(cfg.model) if cfg.model else None
+    model_is_lora = model_path and model_path.is_dir() and (model_path / "adapter_config.json").exists()
+    use_lora = cfg.training.training_type.lower() == "lora"
+
+    if model_is_lora and not use_lora:
+        typer.echo(
+            "Error: Cannot do full finetuning on a LoRA adapter. "
+            "Use --training-type lora or provide a base model.",
+            err=True,
+        )
+        raise typer.Exit(code=2)
+
+    from backend.trainer import UnslothTrainer
+
+    trainer = UnslothTrainer()
+
+    # Load model (trainer.is_vlm is set after this)
+    if not trainer.load_model(
+        model_name=cfg.model,
+        max_seq_length=cfg.training.max_seq_length,
+        load_in_4bit=cfg.training.load_in_4bit if use_lora else False,
+        hf_token=hf_token,
+    ):
+        typer.echo("Model load failed", err=True)
+        raise typer.Exit(code=1)
+
+    is_vision = trainer.is_vlm
+
+    if not trainer.prepare_model_for_training(**cfg.model_kwargs(use_lora, is_vision)):
+        typer.echo("Model preparation failed", err=True)
+        raise typer.Exit(code=1)
+
+    ds = trainer.load_and_format_dataset(
+        dataset_source=cfg.data.dataset or "",
+        format_type=cfg.data.format_type,
+        local_datasets=cfg.data.local_dataset,
+    )
+    if ds is None:
+        typer.echo("Dataset load failed", err=True)
+        raise typer.Exit(code=1)
+
+    training_kwargs = cfg.training_kwargs()
+    training_kwargs["wandb_token"] = wandb_token  # CLI/env takes precedence
+    started = trainer.start_training(dataset=ds, **training_kwargs)
+
+    if not started:
+        typer.echo("Training failed to start", err=True)
+        raise typer.Exit(code=1)
+
+    try:
+        while trainer.training_thread and trainer.training_thread.is_alive():
+            time.sleep(1)
+    except KeyboardInterrupt:
+        typer.echo("Stopping training (Ctrl+C detected)...")
+        trainer.stop_training()
+    finally:
+        if trainer.training_thread:
+            trainer.training_thread.join()
+
+    final = trainer.get_training_progress()
+    if getattr(final, "error", None):
+        typer.echo(f"Training error: {final.error}", err=True)
+        raise typer.Exit(code=1)

cli/commands/ui.py

@@ -0,0 +1,19 @@
+import typer
+
+
+def ui(
+    port: int = typer.Option(8000, "--port", "-p", help="Port to run the UI server on."),
+    host: str = typer.Option("0.0.0.0", "--host", "-H", help="Host address to bind to."),
+    share: bool = typer.Option(True, "--share", "-s", help="Create a public Gradio share link."),
+):
+    """Launch the Unsloth web UI for training, inference, and export."""
+    from app import demo, script_dir
+
+    typer.echo(f"Starting Unsloth UI on http://{host}:{port}")
+
+    demo.launch(
+        share=share,
+        server_port=port,
+        server_name=host,
+        favicon_path=f"{script_dir}/assets/favicon-32x32.png",
+    )

cli/config.py

@@ -0,0 +1,141 @@
+from pathlib import Path
+from typing import Literal, Optional, List
+
+import yaml
+from pydantic import BaseModel, Field
+
+
+class DataConfig(BaseModel):
+    dataset: Optional[str] = None
+    local_dataset: Optional[List[str]] = None
+    format_type: Literal["auto", "alpaca", "chatml", "sharegpt"] = "auto"
+
+
+class TrainingConfig(BaseModel):
+    training_type: Literal["lora", "full"] = "lora"
+    max_seq_length: int = 2048
+    load_in_4bit: bool = True
+    output_dir: Path = Path("./outputs")
+    num_epochs: int = 3
+    learning_rate: float = 2e-4
+    batch_size: int = 2
+    gradient_accumulation_steps: int = 4
+    warmup_steps: int = 5
+    max_steps: int = 0
+    save_steps: int = 0
+    weight_decay: float = 0.01
+    random_seed: int = 3407
+    packing: bool = False
+    train_on_completions: bool = False
+    gradient_checkpointing: Literal["unsloth", "true", "none"] = "unsloth"
+
+
+class LoraConfig(BaseModel):
+    lora_r: int = 64
+    lora_alpha: int = 16
+    lora_dropout: float = 0.0
+    target_modules: str = "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj"
+    vision_all_linear: bool = False
+    use_rslora: bool = False
+    use_loftq: bool = False
+    finetune_vision_layers: bool = True
+    finetune_language_layers: bool = True
+    finetune_attention_modules: bool = True
+    finetune_mlp_modules: bool = True
+
+
+class LoggingConfig(BaseModel):
+    enable_wandb: bool = False
+    wandb_project: str = "unsloth-training"
+    wandb_token: Optional[str] = None
+    enable_tensorboard: bool = False
+    tensorboard_dir: str = "runs"
+    hf_token: Optional[str] = None
+
+
+class Config(BaseModel):
+    model: Optional[str] = None
+    data: DataConfig = Field(default_factory=DataConfig)
+    training: TrainingConfig = Field(default_factory=TrainingConfig)
+    lora: LoraConfig = Field(default_factory=LoraConfig)
+    logging: LoggingConfig = Field(default_factory=LoggingConfig)
+
+    def apply_overrides(self, **kwargs):
+        """Apply CLI overrides by matching arg names to config fields."""
+        for key, value in kwargs.items():
+            if value is None:
+                continue
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                for section in (self.data, self.training, self.lora, self.logging):
+                    if hasattr(section, key):
+                        setattr(section, key, value)
+                        break
+
+    def model_kwargs(self, use_lora: bool, is_vision: bool) -> dict:
+        """Return kwargs for trainer.prepare_model_for_training()."""
+        # Determine target modules based on model type
+        if use_lora and is_vision:
+            # Vision models expect a string (e.g., "all-linear"); fall back to None to use trainer defaults
+            target_modules = "all-linear" if self.lora.vision_all_linear else None
+        else:
+            parsed = [m.strip() for m in str(self.lora.target_modules).split(",") if m and m.strip()]
+            target_modules = parsed or None
+
+        return {
+            "use_lora": use_lora,
+            "finetune_vision_layers": self.lora.finetune_vision_layers,
+            "finetune_language_layers": self.lora.finetune_language_layers,
+            "finetune_attention_modules": self.lora.finetune_attention_modules,
+            "finetune_mlp_modules": self.lora.finetune_mlp_modules,
+            "target_modules": target_modules,
+            "lora_r": self.lora.lora_r,
+            "lora_alpha": self.lora.lora_alpha,
+            "lora_dropout": self.lora.lora_dropout,
+            "use_gradient_checkpointing": self.training.gradient_checkpointing,
+            "use_rslora": self.lora.use_rslora,
+            "use_loftq": self.lora.use_loftq,
+        }
+
+    def training_kwargs(self) -> dict:
+        """Return kwargs for trainer.start_training()."""
+        return {
+            "output_dir": str(self.training.output_dir),
+            "num_epochs": self.training.num_epochs,
+            "learning_rate": self.training.learning_rate,
+            "batch_size": self.training.batch_size,
+            "gradient_accumulation_steps": self.training.gradient_accumulation_steps,
+            "warmup_steps": self.training.warmup_steps,
+            "max_steps": self.training.max_steps,
+            "save_steps": self.training.save_steps,
+            "weight_decay": self.training.weight_decay,
+            "random_seed": self.training.random_seed,
+            "packing": self.training.packing,
+            "train_on_completions": self.training.train_on_completions,
+            "max_seq_length": self.training.max_seq_length,
+            "enable_wandb": self.logging.enable_wandb,
+            "wandb_project": self.logging.wandb_project,
+            "wandb_token": self.logging.wandb_token,
+            "enable_tensorboard": self.logging.enable_tensorboard,
+            "tensorboard_dir": self.logging.tensorboard_dir,
+        }
+
+
+def load_config(path: Optional[Path]) -> Config:
+    """Load config from YAML/JSON file, or return defaults if no path given."""
+    if not path:
+        return Config()
+
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Config file not found: {path}")
+
+    text = path.read_text(encoding="utf-8")
+    if path.suffix.lower() in {".yaml", ".yml"}:
+        data = yaml.safe_load(text) or {}
+    else:
+        import json
+        data = json.loads(text or "{}")
+
+    return Config(**data)

cli/options.py

@@ -0,0 +1,148 @@
+"""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())
+        original_param_names = {p.name for p in original_params}
+
+        # Build new parameters: config fields first, then original params
+        new_params = []
+
+        for field_name, field_info in fields:
+            # Skip fields already defined in function signature (e.g., with envvar)
+            if field_name in original_param_names:
+                continue
+            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]
+                    # Only delete if not an explicitly declared parameter
+                    if key not in original_param_names:
+                        del kwargs[key]
+
+            kwargs["config_overrides"] = config_overrides
+            return func(*args, **kwargs)
+
+        wrapper.__signature__ = new_sig
+        return wrapper
+
+    return decorator