import logging import sys import time from pathlib import Path from typing import Optional, List import typer from cli.config import Config, load_config app = typer.Typer( help="Command-line interface for Unsloth training, chat, and export.", context_settings={"help_option_names": ["-h", "--help"]}, ) def configure_logging(verbose: bool): level = logging.DEBUG if verbose else logging.INFO logging.basicConfig( level=level, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s", datefmt="%H:%M:%S", ) @app.command() 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.", ), verbose: bool = typer.Option(False, "--verbose/--quiet"), ): """ 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) # Apply CLI overrides cli_args = {k: v for k, v in locals().items() if k not in ("config", "verbose", "hf_token", "cfg")} cfg.apply_overrides(**cli_args) # Validate required fields 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) # Lazy imports to avoid triggering Unsloth patches on --help from backend.trainer import UnslothTrainer from backend.model_config import ModelConfig configure_logging(verbose) trainer = UnslothTrainer() model_config = ModelConfig.from_ui_selection( dropdown_value=cfg.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) is_vision = model_config.is_vision use_lora = cfg.training.training_type.lower() == "lora" if not trainer.load_model( model_name=model_config.identifier, 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) 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) started = trainer.start_training(dataset=ds, **cfg.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) @app.command() 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"), verbose: bool = typer.Option(False, "--verbose/--quiet"), ): """ Run a single inference using the specified model. """ # Lazy imports to avoid triggering Unsloth patches on --help from backend.model_config import ModelConfig from backend.inference import get_inference_backend configure_logging(verbose) 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: # Backend yields cumulative text; print only the delta delta = chunk[len(previous):] if delta: sys.stdout.write(delta) sys.stdout.flush() previous = chunk sys.stdout.write("\n") sys.stdout.flush() @app.command("list-checkpoints") 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}") if __name__ == "__main__": app()