add config support + example configs, etc.
This commit is contained in:
Dan Saunders
parent
42490cfbc4
commit
4ef25032c1
1 changed files with 306 additions and 132 deletions
438
cli.py
438
cli.py
|
|
@ -1,18 +1,16 @@
|
|||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, TYPE_CHECKING
|
||||
from typing import Optional, List
|
||||
|
||||
import typer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# Only import for type hints to avoid triggering heavy backend initialization on CLI --help
|
||||
from backend.trainer import TrainingProgress
|
||||
import yaml
|
||||
|
||||
app = typer.Typer(
|
||||
help="Command-line interface for Unsloth training, chat, and export.",
|
||||
context_settings={"help_option_names": ["-h", "--help"]},
|
||||
context_settings={"help_option_names": ["-h", "--help"]},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -25,33 +23,90 @@ def configure_logging(verbose: bool):
|
|||
)
|
||||
|
||||
|
||||
def _print_progress(progress):
|
||||
parts = []
|
||||
if progress.step:
|
||||
if progress.total_steps:
|
||||
parts.append(f"step {progress.step}/{progress.total_steps}")
|
||||
else:
|
||||
parts.append(f"step {progress.step}")
|
||||
if progress.epoch:
|
||||
parts.append(f"epoch {progress.epoch}")
|
||||
if progress.loss:
|
||||
parts.append(f"loss {progress.loss:.4f}")
|
||||
if progress.learning_rate:
|
||||
parts.append(f"lr {progress.learning_rate:.2e}")
|
||||
status = progress.status_message or ""
|
||||
if not parts and status:
|
||||
line = status
|
||||
def _load_config(config_path: Optional[Path]) -> dict:
|
||||
if not config_path:
|
||||
return {}
|
||||
path = Path(config_path)
|
||||
if not path.exists():
|
||||
raise typer.BadParameter(f"Config file not found: {config_path}")
|
||||
text = path.read_text(encoding="utf-8")
|
||||
if path.suffix.lower() in {".yaml", ".yml"}:
|
||||
return yaml.safe_load(text) or {}
|
||||
else:
|
||||
line = " | ".join(parts)
|
||||
if status:
|
||||
line = f"{line} | {status}"
|
||||
if line:
|
||||
typer.echo(line)
|
||||
return json.loads(text or "{}")
|
||||
|
||||
|
||||
def _flatten_config(cfg: dict) -> dict:
|
||||
"""
|
||||
Flatten nested config sections into a single dict.
|
||||
|
||||
Expected sections:
|
||||
data: dataset, local_dataset, format_type
|
||||
training: training_type, max_seq_length, load_in_4bit, output_dir, etc.
|
||||
lora: lora_r, lora_alpha, lora_dropout, target_modules, etc.
|
||||
vision: finetune_vision_layers, finetune_language_layers, etc.
|
||||
logging: enable_wandb, wandb_project, wandb_token, enable_tensorboard, etc.
|
||||
"""
|
||||
if not isinstance(cfg, dict):
|
||||
return {}
|
||||
|
||||
flattened = {}
|
||||
|
||||
# Handle top-level 'model' key
|
||||
if "model" in cfg:
|
||||
flattened["model"] = cfg["model"]
|
||||
|
||||
sections = ["data", "training", "lora", "vision", "logging"]
|
||||
|
||||
for section in sections:
|
||||
if section in cfg and isinstance(cfg[section], dict):
|
||||
flattened.update(cfg[section])
|
||||
|
||||
return flattened
|
||||
|
||||
|
||||
def _merge_config(cfg: dict, defaults: dict, overrides: dict) -> dict:
|
||||
"""
|
||||
Merge CLI overrides with config and defaults.
|
||||
CLI override wins, then config value, then default.
|
||||
"""
|
||||
merged = {}
|
||||
for key, default in defaults.items():
|
||||
cli_val = overrides.get(key, None)
|
||||
if cli_val is not None:
|
||||
merged[key] = cli_val
|
||||
elif key in cfg and cfg[key] is not None:
|
||||
merged[key] = cfg[key]
|
||||
else:
|
||||
merged[key] = default
|
||||
return merged
|
||||
|
||||
|
||||
@app.command()
|
||||
def train(
|
||||
model: str = typer.Argument(..., help="HF model id or local path."),
|
||||
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",
|
||||
|
|
@ -63,68 +118,182 @@ def train(
|
|||
"--local-dataset",
|
||||
help="Filename(s) under datasets/ to use (e.g. 'alpaca_unsloth.json').",
|
||||
),
|
||||
output_dir: Path = typer.Option(
|
||||
Path("./outputs"),
|
||||
"--output-dir",
|
||||
help="Where to store checkpoints.",
|
||||
format_type: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--format-type",
|
||||
help="Dataset formatting: auto|alpaca|chatml|sharegpt. Defaults to auto.",
|
||||
),
|
||||
training_type: str = typer.Option(
|
||||
"lora",
|
||||
"--training-type",
|
||||
help="Training mode: 'lora' (LoRA/QLoRA) or 'full'.",
|
||||
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."
|
||||
),
|
||||
hf_token: Optional[str] = typer.Option(
|
||||
None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed."
|
||||
save_steps: Optional[int] = typer.Option(
|
||||
None, "--save-steps", help="0 uses trainer defaults."
|
||||
),
|
||||
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
|
||||
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
|
||||
num_epochs: int = typer.Option(3, "--epochs"),
|
||||
learning_rate: float = typer.Option(2e-4, "--lr"),
|
||||
batch_size: int = typer.Option(2, "--batch-size"),
|
||||
gradient_accumulation_steps: int = typer.Option(4, "--grad-accum"),
|
||||
warmup_steps: int = typer.Option(5, "--warmup-steps"),
|
||||
max_steps: int = typer.Option(0, "--max-steps", help="Overrides epochs if >0."),
|
||||
save_steps: int = typer.Option(0, "--save-steps", help="0 uses trainer defaults."),
|
||||
weight_decay: float = typer.Option(0.01, "--weight-decay"),
|
||||
random_seed: int = typer.Option(3407, "--seed"),
|
||||
packing: bool = typer.Option(False, "--packing/--no-packing"),
|
||||
train_on_completions: bool = typer.Option(
|
||||
False, "--train-on-completions", help="Train on responses only when supported."
|
||||
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: int = typer.Option(64, "--lora-r"),
|
||||
lora_alpha: int = typer.Option(16, "--lora-alpha"),
|
||||
lora_dropout: float = typer.Option(0.0, "--lora-dropout"),
|
||||
gradient_checkpointing: bool = typer.Option(
|
||||
True, "--gradient-checkpointing/--no-gradient-checkpointing"
|
||||
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: str = typer.Option(
|
||||
"q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
|
||||
target_modules: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--target-modules",
|
||||
help="Comma-separated target modules for LoRA.",
|
||||
),
|
||||
use_rslora: bool = typer.Option(False, "--rslora/--no-rslora"),
|
||||
use_loftq: bool = typer.Option(False, "--loftq/--no-loftq"),
|
||||
enable_wandb: bool = typer.Option(False, "--wandb/--no-wandb"),
|
||||
wandb_project: str = typer.Option("unsloth-training", "--wandb-project"),
|
||||
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: bool = typer.Option(
|
||||
False, "--tensorboard/--no-tensorboard", help="Enable TensorBoard logging."
|
||||
enable_tensorboard: Optional[bool] = typer.Option(
|
||||
None, "--tensorboard/--no-tensorboard", help="Enable TensorBoard logging."
|
||||
),
|
||||
tensorboard_dir: str = typer.Option("runs", "--tensorboard-dir"),
|
||||
format_type: str = typer.Option(
|
||||
"auto",
|
||||
"--format-type",
|
||||
help="Dataset formatting: auto|alpaca|chatml|sharegpt.",
|
||||
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.
|
||||
"""
|
||||
if not dataset and not local_dataset:
|
||||
typer.echo("Error: provide --dataset or --local-dataset", err=True)
|
||||
cfg = _load_config(config)
|
||||
cfg = _flatten_config(cfg)
|
||||
|
||||
# Defaults (match previous behavior)
|
||||
defaults = {
|
||||
"model": None,
|
||||
"training_type": "lora",
|
||||
"max_seq_length": 2048,
|
||||
"load_in_4bit": True,
|
||||
"output_dir": Path("./outputs"),
|
||||
"dataset": None,
|
||||
"local_dataset": None,
|
||||
"format_type": "auto",
|
||||
"num_epochs": 3,
|
||||
"learning_rate": 2e-4,
|
||||
"batch_size": 2,
|
||||
"gradient_accumulation_steps": 4,
|
||||
"warmup_steps": 5,
|
||||
"max_steps": 0,
|
||||
"save_steps": 0,
|
||||
"weight_decay": 0.01,
|
||||
"random_seed": 3407,
|
||||
"packing": False,
|
||||
"train_on_completions": False,
|
||||
"lora_r": 64,
|
||||
"lora_alpha": 16,
|
||||
"lora_dropout": 0.0,
|
||||
"gradient_checkpointing": True,
|
||||
"target_modules": "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
|
||||
"vision_all_linear": False,
|
||||
"finetune_vision_layers": True,
|
||||
"finetune_language_layers": True,
|
||||
"finetune_attention_modules": True,
|
||||
"finetune_mlp_modules": True,
|
||||
"use_rslora": False,
|
||||
"use_loftq": False,
|
||||
"enable_wandb": False,
|
||||
"wandb_project": "unsloth-training",
|
||||
"enable_tensorboard": False,
|
||||
"tensorboard_dir": "runs",
|
||||
}
|
||||
|
||||
overrides = {
|
||||
"training_type": training_type,
|
||||
"max_seq_length": max_seq_length,
|
||||
"load_in_4bit": load_in_4bit,
|
||||
"output_dir": output_dir,
|
||||
"dataset": dataset,
|
||||
"local_dataset": local_dataset,
|
||||
"format_type": format_type,
|
||||
"num_epochs": num_epochs,
|
||||
"learning_rate": learning_rate,
|
||||
"batch_size": batch_size,
|
||||
"gradient_accumulation_steps": gradient_accumulation_steps,
|
||||
"warmup_steps": warmup_steps,
|
||||
"max_steps": max_steps,
|
||||
"save_steps": save_steps,
|
||||
"weight_decay": weight_decay,
|
||||
"random_seed": random_seed,
|
||||
"packing": packing,
|
||||
"train_on_completions": train_on_completions,
|
||||
"lora_r": lora_r,
|
||||
"lora_alpha": lora_alpha,
|
||||
"lora_dropout": lora_dropout,
|
||||
"gradient_checkpointing": gradient_checkpointing,
|
||||
"target_modules": target_modules,
|
||||
"vision_all_linear": vision_all_linear,
|
||||
"finetune_vision_layers": finetune_vision_layers,
|
||||
"finetune_language_layers": finetune_language_layers,
|
||||
"finetune_attention_modules": finetune_attention_modules,
|
||||
"finetune_mlp_modules": finetune_mlp_modules,
|
||||
"use_rslora": use_rslora,
|
||||
"use_loftq": use_loftq,
|
||||
"enable_wandb": enable_wandb,
|
||||
"wandb_project": wandb_project,
|
||||
"wandb_token": wandb_token,
|
||||
"enable_tensorboard": enable_tensorboard,
|
||||
"tensorboard_dir": tensorboard_dir,
|
||||
}
|
||||
|
||||
merged = _merge_config(cfg, defaults, overrides)
|
||||
|
||||
model_val = merged.get("model")
|
||||
if not model_val:
|
||||
typer.echo("Error: provide --model or set model in --config", err=True)
|
||||
raise typer.Exit(code=2)
|
||||
|
||||
# Convert specific types
|
||||
output_dir_val = Path(merged["output_dir"])
|
||||
dataset_val = merged.get("dataset")
|
||||
local_dataset_val = merged.get("local_dataset")
|
||||
|
||||
if not dataset_val and not local_dataset_val:
|
||||
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
|
||||
|
|
@ -134,82 +303,86 @@ def train(
|
|||
configure_logging(verbose)
|
||||
trainer = UnslothTrainer()
|
||||
|
||||
def progress_cb(progress: "TrainingProgress"):
|
||||
_print_progress(progress)
|
||||
|
||||
trainer.add_progress_callback(progress_cb)
|
||||
|
||||
model_config = ModelConfig.from_ui_selection(
|
||||
dropdown_value=model, search_value=None, hf_token=hf_token, is_lora=False
|
||||
dropdown_value=model_val, 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)
|
||||
|
||||
typer.echo(f"Loading model: {model_config.identifier}")
|
||||
is_vision = model_config.is_vision
|
||||
|
||||
if not trainer.load_model(
|
||||
model_name=model_config.identifier,
|
||||
max_seq_length=max_seq_length,
|
||||
load_in_4bit=load_in_4bit if training_type.lower() == "lora" else False,
|
||||
max_seq_length=merged["max_seq_length"],
|
||||
load_in_4bit=merged["load_in_4bit"]
|
||||
if merged["training_type"].lower() == "lora"
|
||||
else False,
|
||||
hf_token=hf_token,
|
||||
):
|
||||
typer.echo("Model load failed", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
use_lora = training_type.lower() == "lora"
|
||||
typer.echo(f"Preparing model for {'LoRA' if use_lora else 'full'} finetuning...")
|
||||
use_lora = merged["training_type"].lower() == "lora"
|
||||
|
||||
# Match UI behavior for target modules:
|
||||
# - Text: use parsed target modules list
|
||||
# - Vision: if vision_all_linear, use ["all-linear"]; otherwise empty list
|
||||
target_modules_list = [
|
||||
m.strip() for m in merged["target_modules"].split(",") if m.strip()
|
||||
]
|
||||
if use_lora and is_vision:
|
||||
if merged["vision_all_linear"]:
|
||||
target_modules_list = ["all-linear"]
|
||||
else:
|
||||
target_modules_list = []
|
||||
|
||||
if not trainer.prepare_model_for_training(
|
||||
use_lora=use_lora,
|
||||
finetune_vision_layers=True,
|
||||
finetune_language_layers=True,
|
||||
finetune_attention_modules=True,
|
||||
finetune_mlp_modules=True,
|
||||
target_modules=[m.strip() for m in target_modules.split(",") if m.strip()],
|
||||
lora_r=lora_r,
|
||||
lora_alpha=lora_alpha,
|
||||
lora_dropout=lora_dropout,
|
||||
use_gradient_checkpointing=gradient_checkpointing,
|
||||
use_rslora=use_rslora,
|
||||
use_loftq=use_loftq,
|
||||
finetune_vision_layers=merged["finetune_vision_layers"],
|
||||
finetune_language_layers=merged["finetune_language_layers"],
|
||||
finetune_attention_modules=merged["finetune_attention_modules"],
|
||||
finetune_mlp_modules=merged["finetune_mlp_modules"],
|
||||
target_modules=target_modules_list,
|
||||
lora_r=merged["lora_r"],
|
||||
lora_alpha=merged["lora_alpha"],
|
||||
lora_dropout=merged["lora_dropout"],
|
||||
use_gradient_checkpointing=merged["gradient_checkpointing"],
|
||||
use_rslora=merged["use_rslora"],
|
||||
use_loftq=merged["use_loftq"],
|
||||
):
|
||||
typer.echo("Model preparation failed", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
if not dataset and not local_dataset:
|
||||
typer.echo("Provide --dataset or --local-dataset", err=True)
|
||||
raise typer.Exit(code=2)
|
||||
|
||||
typer.echo("Loading dataset...")
|
||||
ds = trainer.load_and_format_dataset(
|
||||
dataset_source=dataset or "",
|
||||
format_type=format_type,
|
||||
local_datasets=local_dataset,
|
||||
dataset_source=dataset_val or "",
|
||||
format_type=merged["format_type"],
|
||||
local_datasets=local_dataset_val,
|
||||
)
|
||||
if ds is None:
|
||||
typer.echo("Dataset load failed", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
typer.echo("Starting training...")
|
||||
started = trainer.start_training(
|
||||
dataset=ds,
|
||||
output_dir=str(output_dir),
|
||||
num_epochs=num_epochs,
|
||||
learning_rate=learning_rate,
|
||||
batch_size=batch_size,
|
||||
gradient_accumulation_steps=gradient_accumulation_steps,
|
||||
warmup_steps=warmup_steps,
|
||||
max_steps=max_steps,
|
||||
save_steps=save_steps,
|
||||
weight_decay=weight_decay,
|
||||
random_seed=random_seed,
|
||||
packing=packing,
|
||||
train_on_completions=train_on_completions,
|
||||
enable_wandb=enable_wandb,
|
||||
wandb_project=wandb_project,
|
||||
wandb_token=wandb_token,
|
||||
enable_tensorboard=enable_tensorboard,
|
||||
tensorboard_dir=tensorboard_dir,
|
||||
max_seq_length=max_seq_length,
|
||||
output_dir=str(output_dir_val),
|
||||
num_epochs=merged["num_epochs"],
|
||||
learning_rate=merged["learning_rate"],
|
||||
batch_size=merged["batch_size"],
|
||||
gradient_accumulation_steps=merged["gradient_accumulation_steps"],
|
||||
warmup_steps=merged["warmup_steps"],
|
||||
max_steps=merged["max_steps"],
|
||||
save_steps=merged["save_steps"],
|
||||
weight_decay=merged["weight_decay"],
|
||||
random_seed=merged["random_seed"],
|
||||
packing=merged["packing"],
|
||||
train_on_completions=merged["train_on_completions"],
|
||||
enable_wandb=merged["enable_wandb"],
|
||||
wandb_project=merged["wandb_project"],
|
||||
wandb_token=merged.get("wandb_token"),
|
||||
enable_tensorboard=merged["enable_tensorboard"],
|
||||
tensorboard_dir=merged["tensorboard_dir"],
|
||||
max_seq_length=merged["max_seq_length"],
|
||||
)
|
||||
|
||||
if not started:
|
||||
|
|
@ -218,9 +391,7 @@ def train(
|
|||
|
||||
try:
|
||||
while trainer.training_thread and trainer.training_thread.is_alive():
|
||||
progress = trainer.get_training_progress()
|
||||
_print_progress(progress)
|
||||
time.sleep(5)
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
typer.echo("Stopping training (Ctrl+C detected)...")
|
||||
trainer.stop_training()
|
||||
|
|
@ -229,16 +400,15 @@ def train(
|
|||
trainer.training_thread.join()
|
||||
|
||||
final = trainer.get_training_progress()
|
||||
if final.error:
|
||||
if getattr(final, "error", None):
|
||||
typer.echo(f"Training error: {final.error}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
typer.echo(final.status_message or "Training complete")
|
||||
|
||||
|
||||
@app.command()
|
||||
def chat(
|
||||
def inference(
|
||||
model: str = typer.Argument(..., help="HF model id or local path."),
|
||||
prompt: str = typer.Argument(..., help="User prompt to send."),
|
||||
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."
|
||||
),
|
||||
|
|
@ -257,7 +427,7 @@ def chat(
|
|||
verbose: bool = typer.Option(False, "--verbose/--quiet"),
|
||||
):
|
||||
"""
|
||||
Run a single chat turn using the inference backend.
|
||||
Run a single inference using the specified model.
|
||||
"""
|
||||
# Lazy imports to avoid triggering Unsloth patches on --help
|
||||
from backend.model_config import ModelConfig
|
||||
|
|
@ -272,7 +442,6 @@ def chat(
|
|||
typer.echo("Could not resolve model config", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
typer.echo(f"Loading model: {model_config.identifier}")
|
||||
if not inference_backend.load_model(
|
||||
config=model_config,
|
||||
max_seq_length=max_seq_length,
|
||||
|
|
@ -294,9 +463,14 @@ def chat(
|
|||
)
|
||||
|
||||
typer.echo("Assistant:", nl=True)
|
||||
previous = ""
|
||||
for chunk in stream:
|
||||
sys.stdout.write(chunk)
|
||||
sys.stdout.flush()
|
||||
# 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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue