316 lines
11 KiB
Python
316 lines
11 KiB
Python
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()
|