Final cleanup

This commit is contained in:
Roland Tannous 2026-03-12 18:28:04 +00:00
commit 985d2e43ee
123 changed files with 7474 additions and 5805 deletions

View file

@ -13,14 +13,14 @@ 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."
Path("./outputs"), "--outputs-dir", help = "Directory that holds training runs."
),
):
"""List checkpoints detected in the outputs directory."""
from studio.backend.core.export import ExportBackend
backend = ExportBackend()
checkpoints = backend.scan_checkpoints(outputs_dir=str(outputs_dir))
checkpoints = backend.scan_checkpoints(outputs_dir = str(outputs_dir))
if not checkpoints:
typer.echo("No checkpoints found.")
raise typer.Exit()
@ -33,96 +33,100 @@ def list_checkpoints(
def export(
checkpoint: Path = typer.Argument(..., help="Path to checkpoint directory."),
output_dir: Path = typer.Argument(..., help="Directory to save exported model."),
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)}",
help = f"Export format: {', '.join(EXPORT_FORMATS)}",
),
quantization: str = typer.Option(
"q4_k_m",
"--quantization",
"-q",
help=f"GGUF quantization method: {', '.join(GGUF_QUANTS)}",
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."
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)."
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."
None, "--hf-token", envvar = "HF_TOKEN", help = "HuggingFace token."
),
private: bool = typer.Option(
False, "--private", help="Make the HuggingFace repo private."
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)
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)
typer.echo("Error: --repo-id required when using --push-to-hub", err = True)
raise typer.Exit(code = 2)
from studio.backend.core.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,
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(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,
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,
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,
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,
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(f"Error: {message}", err = True)
raise typer.Exit(code = 1)
typer.echo(message)