zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/cli/commands.py

@@ -0,0 +1,459 @@
+"""Non-interactive Typer subcommands.
+
+Each command builds an ephemeral (scratch) context, runs the same intents the
+interactive surfaces use, and renders with the shared Rich renderers — so
+``zeroquantz recommend`` prints exactly what ``/recommend`` shows in the TUI.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import typer
+from rich.console import Console
+
+from zeroquantz.agent.dispatcher import CommandResult, Dispatcher
+from zeroquantz.agent.intents import Intent, IntentKind
+from zeroquantz.core.exceptions import ZeroQuantzError
+
+# AppContext (hardware detection + registries), render_result, and Objective are
+# imported lazily inside the functions that use them. Building the Typer app and
+# registering commands must stay cheap so `zeroquantz --help`/`--version` and shell
+# completion don't drag in the detection + optimization stack.
+if TYPE_CHECKING:
+    from zeroquantz.core.context import AppContext
+
+console = Console()
+
+
+def _context(gpu: str | None = None) -> AppContext:
+    from zeroquantz.core.context import AppContext
+
+    ctx = AppContext.create()  # scratch session, never persisted
+    if gpu and gpu.lower() != "auto":
+        from zeroquantz.hardware.detector import HardwareDetector
+
+        ctx.hardware = HardwareDetector.simulate(gpu)
+        ctx.session.hardware = ctx.hardware
+    return ctx
+
+
+def _cli_progress(stage: str, fraction: float) -> None:
+    console.print(f"[grey50]  {stage} … {fraction * 100:.0f}%[/grey50]")
+
+
+def _emit(result: CommandResult) -> None:
+    from zeroquantz.render import render_result
+
+    console.print(render_result(result))
+    if not result.success:
+        raise typer.Exit(code=1)
+
+
+def _apply_goal(ctx: AppContext, **updates: object) -> None:
+    from zeroquantz.optimization.constraints import Objective
+
+    clean = {k: v for k, v in updates.items() if v is not None}
+    if "objective" in clean and clean["objective"] is not None:
+        clean["objective"] = Objective(str(clean["objective"]))
+    if clean:
+        ctx.session.goal = ctx.session.goal.with_updates(**clean)
+
+
+def register(app: typer.Typer) -> None:
+    disp = Dispatcher()
+
+    @app.command()
+    def hardware() -> None:
+        """Show the detected GPU, VRAM, and precision support."""
+        _emit(disp.dispatch(Intent(IntentKind.HARDWARE), _context()))
+
+    @app.command()
+    def inspect(model: str = typer.Argument(..., help="Hugging Face id or local path.")) -> None:
+        """Inspect a model's architecture and size (metadata only)."""
+        ctx = _context()
+        _emit(disp.dispatch(Intent(IntentKind.LOAD_MODEL, {"model_id": model}), ctx))
+
+    @app.command()
+    def recommend(
+        model: str = typer.Argument(..., help="Hugging Face id or local path."),
+        max_vram: float | None = typer.Option(None, "--max-vram", help="VRAM budget in GB."),
+        max_size: float | None = typer.Option(None, "--max-size", help="Model size budget in GB."),
+        runtime: str | None = typer.Option(None, "--runtime", help="Target runtime (vllm, transformers)."),
+        objective: str = typer.Option("balanced", "--objective", help="quality | speed | memory | balanced."),
+        context: int = typer.Option(4096, "--context", help="Context length (tokens) for the fit check."),
+        gpu: str | None = typer.Option(None, "--gpu", help="Plan for a named GPU (e.g. 'RTX 4090') or 'auto'."),
+        show_all: bool = typer.Option(False, "--all", help="Show every catalogued format, not just the top 15."),
+    ) -> None:
+        """Rank quantization strategies for a model under your constraints."""
+        ctx = _context(gpu)
+        res = disp.dispatch(Intent(IntentKind.LOAD_MODEL, {"model_id": model}), ctx)
+        if not res.success:
+            _emit(res)
+        _apply_goal(
+            ctx, max_vram_gb=max_vram, max_model_size_gb=max_size,
+            runtime=runtime, objective=objective, context_length=context,
+        )
+        result = disp.dispatch(Intent(IntentKind.RECOMMEND), ctx)
+        if result.success and not show_all and isinstance(result.payload, dict):
+            result.payload["ranked"] = result.payload["ranked"][:15]
+        _emit(result)
+
+    @app.command()
+    def plan(
+        model: str = typer.Argument(...),
+        max_vram: float | None = typer.Option(None, "--max-vram", help="VRAM budget in GB."),
+        gpu: str | None = typer.Option(None, "--gpu", help="Plan for a named GPU or 'auto'."),
+    ) -> None:
+        """Build a mixed-precision plan for a memory budget."""
+        ctx = _context(gpu)
+        res = disp.dispatch(Intent(IntentKind.LOAD_MODEL, {"model_id": model}), ctx)
+        if not res.success:
+            _emit(res)
+        _apply_goal(ctx, max_vram_gb=max_vram)
+        _emit(disp.dispatch(Intent(IntentKind.PLAN), ctx))
+
+    @app.command()
+    def profile(
+        model: str = typer.Argument(...),
+        measured: bool = typer.Option(False, "--measured", help="Run a measured pass (needs torch + GPU)."),
+    ) -> None:
+        """Estimate per-layer sensitivity (heuristic, or --measured)."""
+        ctx = _context()
+        res = disp.dispatch(Intent(IntentKind.LOAD_MODEL, {"model_id": model}), ctx)
+        if not res.success:
+            _emit(res)
+        if measured:
+            from zeroquantz.profiling.sensitivity import SensitivityProfiler
+
+            try:
+                sp = SensitivityProfiler.profile(model, progress=_cli_progress)
+                _emit(CommandResult("profile", "Measured layer sensitivity.", sp))
+            except ZeroQuantzError as exc:
+                _emit(CommandResult("error", exc.format(), success=False))
+        else:
+            _emit(disp.dispatch(Intent(IntentKind.PROFILE), ctx))
+
+    @app.command()
+    def quantize(
+        model: str = typer.Argument(...),
+        backend: str | None = typer.Option(None, "--backend", help="Backend name (bitsandbytes, torchao)."),
+        method: str | None = typer.Option(None, "--method", help="Method (nf4, int8_bnb, int8_torchao, int4_torchao)."),
+        output: str | None = typer.Option(None, "--output", "-o", help="Directory to write the quantized model."),
+    ) -> None:
+        """Quantize a model (requires a backend + CUDA GPU)."""
+        ctx = _context()
+        res = disp.dispatch(Intent(IntentKind.LOAD_MODEL, {"model_id": model}), ctx)
+        if not res.success:
+            _emit(res)
+        result = disp.dispatch(Intent(IntentKind.QUANTIZE, {"method": method}), ctx, _cli_progress)
+        from zeroquantz.render import render_result
+
+        console.print(render_result(result))
+        if result.success and output:
+            _emit(disp.dispatch(Intent(IntentKind.EXPORT, {"path": output}), ctx))
+        elif not result.success:
+            raise typer.Exit(code=1)
+
+    @app.command()
+    def benchmark(
+        model: str = typer.Argument(..., help="Model id or a quantized directory."),
+        prompts: str | None = typer.Option(None, "--prompts", help="JSONL/TXT file of prompts."),
+        runs: int = typer.Option(5, "--runs"),
+        warmup: int = typer.Option(2, "--warmup"),
+        quality: bool = typer.Option(False, "--quality", help="Also measure perplexity."),
+    ) -> None:
+        """Benchmark VRAM, latency, and throughput (requires torch + GPU)."""
+        from zeroquantz.benchmark.runner import BenchmarkRunner
+
+        ctx = _context()
+        prompt_list = _load_prompts(prompts)
+        try:
+            result = BenchmarkRunner.run(
+                model, prompts=prompt_list, runs=runs, warmup=warmup,
+                measure_quality=quality, hardware=ctx.hardware, progress=_cli_progress,
+            )
+            _emit(CommandResult("benchmark", f"Benchmarked {model}.", result))
+        except ZeroQuantzError as exc:
+            _emit(CommandResult("error", exc.format(), success=False))
+
+    @app.command()
+    def verify(
+        base: str = typer.Argument(..., help="Baseline (full-precision) model."),
+        quantized: str = typer.Argument(..., help="Quantized model directory."),
+        prompts: str | None = typer.Option(None, "--prompts", help="JSONL/TXT file of prompts."),
+    ) -> None:
+        """Compare a quantized model against its baseline (requires torch + GPU)."""
+        from zeroquantz.verification.report import Verifier
+
+        ctx = _context()
+        try:
+            report = Verifier.verify(
+                base, quantized, prompts=_load_prompts(prompts),
+                hardware=ctx.hardware, progress=_cli_progress,
+            )
+            _emit(CommandResult("verification", f"Verification: {report.status}", report))
+        except ZeroQuantzError as exc:
+            _emit(CommandResult("error", exc.format(), success=False))
+
+    @app.command()
+    def backends() -> None:
+        """List quantization backends and whether they are installed."""
+        from rich.table import Table
+
+        ctx = _context()
+        table = Table(title="Quantization backends")
+        table.add_column("Backend", style="cyan")
+        table.add_column("Available")
+        table.add_column("Methods", style="grey62")
+        for b in ctx.backends.all_backends():
+            ok = "[green]yes[/green]" if b.is_available() else "[grey62]no[/grey62]"
+            table.add_row(b.name, ok, ", ".join(b.methods))
+        console.print(table)
+
+    @app.command()
+    def convert(
+        model: str = typer.Argument(..., help="Local model directory (HF safetensors + config.json)."),
+        output: str = typer.Option(..., "--output", "-o", help="Output .gguf path."),
+        quant: str = typer.Option("q4_0", "--quant", help="q4_0 | q8_0 | f16 | f32."),
+    ) -> None:
+        """Convert an HF model dir to GGUF via the native candle prototype (fast, CPU).
+
+        Uses the ``zeroquantz-convert`` Rust binary when built (Cargo.toml + src/ at
+        the repo root); otherwise prints the build steps and the equivalent
+        llama.cpp commands.
+        """
+        import subprocess
+        from pathlib import Path
+
+        binary = _find_convert_binary()
+        if binary is None:
+            _print_convert_unavailable(model, output, quant)
+            raise typer.Exit(code=1)
+        if not Path(model).expanduser().is_dir():
+            console.print(
+                f"[yellow]'{model}' is not a local directory.[/yellow] Download the model "
+                "first (the TUI download step, or `huggingface-cli download`), then pass "
+                "its snapshot folder."
+            )
+            raise typer.Exit(code=1)
+        cmd = [str(binary), "--model", str(Path(model).expanduser()), "--out", output, "--quant", quant]
+        console.print(f"[grey50]$ {' '.join(cmd)}[/grey50]")
+        raise typer.Exit(code=subprocess.run(cmd).returncode)
+
+    @app.command()
+    def serve(
+        model: str = typer.Argument(..., help="Local model dir or HF id to serve."),
+        runtime: str = typer.Option("vllm", "--runtime", "-r", help="vllm | sglang | tensorrt_llm."),
+        format_id: str | None = typer.Option(None, "--format", help="Catalog format id (adds the right --quantization flag)."),
+        port: int | None = typer.Option(None, "--port"),
+        plan: bool = typer.Option(False, "--plan", help="Show the isolated-env build + serve commands without running."),
+    ) -> None:
+        """Serve a model on vLLM/SGLang/TensorRT-LLM — each in its own isolated env.
+
+        The runtime is provisioned in a dedicated venv under
+        ``~/.zeroquantz/envs/serve/<runtime>/`` and served from it. Use ``--plan`` to
+        preview the exact build + serve commands without installing anything.
+        """
+        from zeroquantz.deploy import can_provision, get_target, launch
+        from zeroquantz.deploy.runtime_env import RuntimeEnvRunner
+        from zeroquantz.quantization import catalog
+
+        target = get_target(runtime)
+        if target is None or not can_provision(target):
+            console.print(
+                f"[yellow]'{runtime}' is not an isolated serving runtime.[/yellow] "
+                "Choose one of: vllm, sglang, tensorrt_llm. For llama.cpp / Ollama / "
+                "Transformers, use the generated deploy assets instead."
+            )
+            raise typer.Exit(code=1)
+        fmt = catalog.get_format(format_id) if format_id else None
+        if format_id and fmt is None:
+            console.print(f"[yellow]Unknown format '{format_id}'.[/yellow] Try `zeroquantz recommend {model}`.")
+            raise typer.Exit(code=1)
+
+        runner = RuntimeEnvRunner()
+        env_plan = runner.plan(target)
+        serve_cmd = target.serve_command(model, fmt, port=port)
+        endpoint = f"http://localhost:{port or target.default_port}/v1" if target.openai_compatible else None
+
+        if plan:
+            console.print(f"[bold cyan]{target.label}[/bold cyan] isolated env → [grey50]{env_plan.env_dir}[/grey50]")
+            for c in env_plan.commands:
+                console.print(f"  [cyan]{c}[/cyan]")
+            console.print(f"  [cyan]{serve_cmd}[/cyan]   [grey50](run inside that env)[/grey50]")
+            state = "[green]built[/green]" if runner.is_ready(target) else "[grey62]not built yet[/grey62]"
+            console.print(f"status: {state}   ·   endpoint (once running): {endpoint or 'n/a'}")
+            return
+
+        cuda = None
+        try:
+            from zeroquantz.hardware.detector import HardwareDetector
+
+            cuda = HardwareDetector.detect().cuda_version
+        except Exception:
+            pass
+        console.print(
+            f"[grey50]Provisioning {target.label} in {env_plan.env_dir} "
+            "(first run installs the runtime; this can be several GB)…[/grey50]"
+        )
+        try:
+            proc = launch(target, serve_cmd, cuda_version=cuda, progress=_cli_progress)
+        except Exception as exc:
+            console.print(f"[red]Could not start {target.label}:[/red] {exc}")
+            raise typer.Exit(code=1) from exc
+        console.print(f"[green]Started {target.label}[/green] — {serve_cmd}")
+        if endpoint:
+            console.print(f"endpoint: {endpoint}")
+        console.print("[grey50](serving in the foreground; Ctrl+C to stop)[/grey50]")
+        raise typer.Exit(code=proc.wait())
+
+    @app.command("gguf-quantize")
+    def gguf_quantize(
+        gguf_in: str = typer.Argument(..., metavar="INPUT.gguf", help="Input f16/f32 GGUF to quantize."),
+        output: str = typer.Option(..., "--output", "-o", help="Output quantized .gguf path."),
+        type_: str = typer.Option("Q4_K_M", "--type", "-t", help="Q4_K_M | Q5_K_M | Q6_K | Q8_0 | Q4_0 …"),
+        verify: bool = typer.Option(True, "--verify/--no-verify", help="Load-test the output after quantizing."),
+    ) -> None:
+        """Quantize a GGUF via llama-cpp-python (pip-only; llama.cpp-loadable output).
+
+        Runs in the isolated ``tools/gguf`` env, built on first use from prebuilt
+        wheels — no compiler. ``convert`` makes an f16 GGUF; this quantizes it.
+        """
+        from pathlib import Path
+
+        from zeroquantz.quantization.gguf_pipeline import GgufPipeline
+
+        if not Path(gguf_in).expanduser().is_file():
+            console.print(
+                f"[yellow]'{gguf_in}' not found.[/yellow] Produce an f16 GGUF first "
+                "(e.g. `zeroquantz build <model> --keep-f16`)."
+            )
+            raise typer.Exit(code=1)
+        pipe = GgufPipeline()
+        try:
+            out = pipe.quantize(str(Path(gguf_in).expanduser()), output, type_, progress=_cli_progress)
+            if verify:
+                pipe.verify(str(out), progress=_cli_progress)
+        except ZeroQuantzError as exc:
+            console.print(f"[red]Quantize failed:[/red] {exc.format()}")
+            raise typer.Exit(code=1) from exc
+        console.print(f"[green]✓ quantized[/green] {out}")
+
+    @app.command()
+    def build(
+        model: str = typer.Argument(..., help="HF id or local model directory."),
+        quant: str = typer.Option("Q4_K_M", "--quant", "-q", help="GGML type: Q4_K_M | Q5_K_M | Q6_K | Q8_0 …"),
+        output: str | None = typer.Option(None, "--out", "-o", help="Output dir (default ./zeroquantz-<name>-<quant>)."),
+        serve_with: str | None = typer.Option(None, "--serve", help="After building, serve with: llamacpp | ollama."),
+        keep_f16: bool = typer.Option(False, "--keep-f16", help="Keep the intermediate f16 GGUF."),
+        plan: bool = typer.Option(False, "--plan", help="Preview the whole pipeline without running."),
+    ) -> None:
+        """Production GGUF pipeline: download → convert → quantize → verify → (serve).
+
+        Pip-only, no compiler: llama.cpp's convert_hf_to_gguf.py (all architectures +
+        tokenizers) for convert, and llama-cpp-python for quantize + a load-test
+        verify, so the output is a proven llama.cpp/Ollama-loadable GGUF. Each stage
+        hard-fails; `--plan` previews it. Runs in the isolated tools/gguf env.
+        """
+        from pathlib import Path
+
+        from zeroquantz.quantization.gguf_pipeline import GgufPipeline
+
+        is_local = Path(model).expanduser().is_dir()
+        name = (Path(model).expanduser().resolve().name if is_local
+                else model.rstrip("/").split("/")[-1]) or "model"
+        out_dir = Path(output).expanduser() if output else Path.cwd() / f"zeroquantz-{name}-{quant}"
+        pipe = GgufPipeline()
+
+        if plan:
+            model_disp = str(Path(model).expanduser()) if is_local else f"<downloaded {model}>"
+            console.print(f"[bold cyan]GGUF pipeline[/bold cyan]  {model} → {out_dir}  ({quant})")
+            console.print("  [grey50]0 download[/grey50]  " + ("local dir (skip)" if is_local else f"download {model}"))
+            for i, st in enumerate(pipe.plan(model_disp, quant, out_dir), start=1):
+                state = "[green]ready[/green]" if st.ready else "[grey62]needs setup[/grey62]"
+                console.print(f"  [grey50]{i} {st.name}[/grey50]  {state}")
+                console.print(f"      [cyan]{' '.join(st.command)}[/cyan]")
+                if not st.ready and st.hint:
+                    console.print(f"      [grey62]↳ {st.hint}[/grey62]")
+            if serve_with:
+                console.print(f"  [grey50]4 serve[/grey50]  via {serve_with} (llama-server / ollama)")
+            return
+
+        if is_local:
+            model_dir = Path(model).expanduser()
+        else:
+            from zeroquantz.models.downloader import download_model, plan_download
+
+            console.print(f"[grey50]downloading {model} …[/grey50]")
+            files, _total = plan_download(model)
+            model_dir = Path(download_model(model, files))
+
+        def prog(stage: str, frac: float) -> None:
+            console.print(f"[grey50]  {stage} … {frac * 100:.0f}%[/grey50]")
+
+        try:
+            arts = pipe.run(str(model_dir), quant, out_dir, progress=prog, keep_f16=keep_f16)
+        except ZeroQuantzError as exc:
+            console.print(f"[red]Pipeline failed:[/red] {exc.format()}")
+            raise typer.Exit(code=1) from exc
+        console.print(f"[green]✓ built[/green] {arts.quantized_gguf}")
+
+        try:
+            from zeroquantz.core.artifacts import default_quantized_registry
+
+            default_quantized_registry().add(
+                str(arts.quantized_gguf), base_model=model,
+                format_id=f"gguf_{quant.lower()}", method=quant,
+            )
+        except Exception:
+            pass
+
+        if serve_with == "llamacpp":
+            import subprocess
+
+            port = 8080
+            cmd = pipe.serve_command(str(arts.quantized_gguf), port)
+            console.print(f"[grey50]serve:[/grey50] [cyan]{' '.join(cmd)}[/cyan]")
+            console.print(
+                f"endpoint: http://localhost:{port}/v1  "
+                "[grey50](llama-cpp-python server; Ctrl+C to stop)[/grey50]"
+            )
+            raise typer.Exit(code=subprocess.run(cmd).returncode)
+        if serve_with:
+            from zeroquantz.deploy import get_target, is_runtime_available, launch
+
+            target = get_target(serve_with)
+            if target is None:
+                console.print(f"[yellow]Unknown serve runtime '{serve_with}'.[/yellow] Try: llamacpp | ollama.")
+                return
+            serve_cmd = target.serve_command(str(arts.quantized_gguf))
+            console.print(f"[grey50]serve:[/grey50] [cyan]{serve_cmd}[/cyan]")
+            if is_runtime_available(target):
+                launch(target, serve_cmd)
+                console.print(f"[green]started {target.label}[/green]")
+            else:
+                console.print(f"[grey62]{target.label} not installed here — run the command above.[/grey62]")
+
+
+def _find_convert_binary():
+    """Locate the native Rust converter (on PATH or ./target/release)."""
+    from zeroquantz.quantization.native import find_native_binary
+
+    return find_native_binary("zeroquantz-convert")
+
+
+def _print_convert_unavailable(model: str, output: str, quant: str) -> None:
+    console.print("[yellow]Native converter not built.[/yellow] Build it (needs a Rust toolchain):")
+    console.print("  [cyan]cargo build --release[/cyan]  (from the repo root)")
+    console.print("Then re-run `zeroquantz convert`. Meanwhile, the llama.cpp path:")
+    console.print(f"  [cyan]python convert_hf_to_gguf.py {model} --outfile model.f16.gguf[/cyan]")
+    console.print(f"  [cyan]llama-quantize model.f16.gguf {output} {quant.upper()}[/cyan]")
+
+
+def _load_prompts(path: str | None) -> list[str] | None:
+    if not path:
+        return None
+    from zeroquantz.profiling.calibration import CalibrationDataset
+
+    return CalibrationDataset.from_file(path).prompts

zeroquantz/cli/interactive.py

@@ -0,0 +1,56 @@
+"""A plain-text REPL used when the Textual TUI can't run (no TTY, ``--no-tui``).
+
+Same parser, same dispatcher, same renderers as the TUI — just a readline loop
+instead of a full-screen app.
+"""
+
+from __future__ import annotations
+
+from rich.console import Console
+from rich.text import Text
+
+from zeroquantz.agent.dispatcher import Dispatcher
+from zeroquantz.agent.intents import Intent, IntentKind
+from zeroquantz.agent.parser import IntentParser
+from zeroquantz.core.context import AppContext
+from zeroquantz.render import render_result
+from zeroquantz.tui.banner import banner_text
+
+
+def run_repl(*, session_name: str | None = None, continue_last: bool = False) -> None:
+    console = Console()
+    ctx = AppContext.create(session_name=session_name, continue_last=continue_last)
+    dispatcher = Dispatcher()
+
+    console.print(Text(banner_text(console.width - 4), style="bold #c9d1d9"))
+    console.print("Interactive Model Optimization Environment", style="grey62")
+    hw = ctx.hardware
+    if hw.cuda_available:
+        console.print(
+            f"[green]●[/green] {hw.gpu_name}  [grey62]VRAM {hw.total_vram_gb:g} GB · "
+            f"CUDA {hw.cuda_version or 'detected'}[/grey62]"
+        )
+    else:
+        console.print("[yellow]⚠ No CUDA GPU detected — planning-only mode[/yellow]")
+    console.print("Type [cyan]/help[/cyan] for commands, [cyan]/exit[/cyan] to quit.\n")
+
+    def progress(stage: str, fraction: float) -> None:
+        console.print(f"[grey50]  {stage} … {fraction * 100:.0f}%[/grey50]")
+
+    while True:
+        try:
+            text = console.input("[bold green]›[/bold green] ").strip()
+        except (EOFError, KeyboardInterrupt):
+            console.print()
+            break
+        if not text:
+            continue
+        intent = IntentParser.parse(text)
+        if intent.kind is IntentKind.EXIT:
+            break
+        result = dispatcher.dispatch(intent, ctx, progress)
+        console.print(render_result(result))
+        if result.should_exit:
+            break
+
+    console.print("Goodbye.", style="grey62")

zeroquantz/core/__init__.py

@@ -0,0 +1,7 @@
+"""Core session/context/event plumbing shared by every ZeroQuantz surface."""
+
+from __future__ import annotations
+
+from zeroquantz.core import exceptions
+
+__all__ = ["exceptions"]