zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/runtimes/transformers.py

@@ -0,0 +1,53 @@
+"""transformers runtime adapter — the reference/default deployment target."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zeroquantz.runtimes.base import RuntimeBackend, ValidationResult
+
+if TYPE_CHECKING:
+    from zeroquantz.quantization.config import QuantizationConfig
+
+
+class TransformersRuntime(RuntimeBackend):
+    name = "transformers"
+    description = "Hugging Face Transformers (AutoModelForCausalLM)."
+
+    def validate_model(
+        self, path: str, config: QuantizationConfig | None = None
+    ) -> ValidationResult:
+        p = Path(path)
+        messages: list[str] = []
+        warnings: list[str] = []
+        ok = True
+
+        if p.exists() and p.is_dir():
+            if not (p / "config.json").exists():
+                ok = False
+                messages.append("No config.json found in the model directory.")
+        else:
+            messages.append(f"Treating '{path}' as a Hub id (not validated offline).")
+
+        compat = self.compatibility(config.method) if config else None
+        if config and compat is not None:
+            messages.append(f"{config.method}: {compat.label} with transformers.")
+
+        return ValidationResult(
+            runtime=self.name,
+            ok=ok,
+            compat=compat or ValidationResult.model_fields["compat"].default,
+            messages=messages,
+            warnings=warnings,
+            serve_command=self.serve_command(path, config),
+        )
+
+    def serve_command(
+        self, path: str, config: QuantizationConfig | None = None
+    ) -> str | None:
+        # transformers is a library, not a server; show the load snippet.
+        return (
+            "python -c \"from transformers import AutoModelForCausalLM; "
+            f"AutoModelForCausalLM.from_pretrained('{path}', device_map='auto')\""
+        )

zeroquantz/runtimes/vllm.py

@@ -0,0 +1,83 @@
+"""vLLM runtime adapter.
+
+For v0.1 this focuses on *compatibility validation and command generation* rather
+than embedding vLLM internals: it tells you whether a quantized model will serve
+well on vLLM and prints the ``vllm serve`` command to run it.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zeroquantz.runtimes.base import RuntimeBackend, ValidationResult
+from zeroquantz.runtimes.compatibility import RuntimeCompat
+
+if TYPE_CHECKING:
+    from zeroquantz.quantization.config import QuantizationConfig
+
+# ZeroQuantz method -> vLLM --quantization argument.
+_VLLM_QUANT_ARG = {
+    "awq": "awq",
+    "gptq": "gptq",
+    "fp8": "fp8",
+    "nf4": "bitsandbytes",
+    "int8_bnb": "bitsandbytes",
+    "int8_torchao": "torchao",
+    "int4_torchao": "torchao",
+}
+
+
+class VLLMRuntime(RuntimeBackend):
+    name = "vllm"
+    description = "vLLM high-throughput inference server."
+
+    def validate_model(
+        self, path: str, config: QuantizationConfig | None = None
+    ) -> ValidationResult:
+        p = Path(path)
+        messages: list[str] = []
+        warnings: list[str] = []
+        ok = True
+
+        if p.exists() and p.is_dir() and not (p / "config.json").exists():
+            ok = False
+            messages.append("No config.json found in the model directory.")
+
+        compat = RuntimeCompat.UNKNOWN
+        if config is not None:
+            compat = self.compatibility(config.method)
+            messages.append(f"{config.method}: {compat.label} on vLLM.")
+            if compat is RuntimeCompat.UNSUPPORTED:
+                ok = False
+                messages.append(
+                    f"vLLM cannot serve '{config.method}'. Prefer AWQ, GPTQ, or FP8."
+                )
+            elif compat is RuntimeCompat.NOT_RECOMMENDED:
+                warnings.append(
+                    f"'{config.method}' loads on vLLM but is slow; AWQ/GPTQ are preferred."
+                )
+            elif compat is RuntimeCompat.LIMITED:
+                warnings.append(
+                    f"'{config.method}' support on vLLM is limited/experimental."
+                )
+
+        return ValidationResult(
+            runtime=self.name,
+            ok=ok,
+            compat=compat,
+            messages=messages,
+            warnings=warnings,
+            serve_command=self.serve_command(path, config),
+        )
+
+    def serve_command(
+        self, path: str, config: QuantizationConfig | None = None
+    ) -> str | None:
+        parts = ["vllm", "serve", path]
+        if config is not None:
+            arg = _VLLM_QUANT_ARG.get(config.method)
+            if arg is not None:
+                parts.append(f"--quantization {arg}")
+        parts.append("--max-model-len 4096")
+        return " ".join(parts)

zeroquantz/tui/__init__.py

@@ -0,0 +1,13 @@
+"""The Textual interactive terminal UI."""
+
+from __future__ import annotations
+
+__all__ = ["ZeroQuantzApp", "run_tui"]
+
+
+def __getattr__(name: str):  # lazy so importing the package never imports textual
+    if name in ("ZeroQuantzApp", "run_tui"):
+        from zeroquantz.tui.app import ZeroQuantzApp, run_tui
+
+        return {"ZeroQuantzApp": ZeroQuantzApp, "run_tui": run_tui}[name]
+    raise AttributeError(name)

zeroquantz/tui/app.py

@@ -0,0 +1,77 @@
+"""The ZeroQuantz Textual application."""
+
+from __future__ import annotations
+
+from textual.app import App
+from textual.binding import Binding
+
+from zeroquantz.agent.dispatcher import Dispatcher
+from zeroquantz.core.context import AppContext
+from zeroquantz.tui.screens.execute import ExecuteScreen
+from zeroquantz.tui.screens.model_select import ModelSelectScreen
+from zeroquantz.tui.screens.plan import PlanScreen
+from zeroquantz.tui.screens.quantize_select import QuantizeSelectScreen
+from zeroquantz.tui.screens.settings import SettingsScreen
+from zeroquantz.tui.screens.welcome import WelcomeScreen
+from zeroquantz.tui.screens.workspace import WorkspaceScreen
+
+
+class ZeroQuantzApp(App):
+    """Interactive TUI: welcome → (guided) model → quantization → plan → execute."""
+
+    CSS_PATH = "styles/app.tcss"
+    TITLE = "ZeroQuantz"
+    BINDINGS = [Binding("ctrl+q", "quit", "Quit", show=True)]
+
+    def __init__(self, ctx: AppContext | None = None) -> None:
+        super().__init__()
+        self.ctx = ctx or AppContext.create()
+        self.dispatcher = Dispatcher()
+        self.pending_command: str | None = None
+        self.selected_candidate = None
+        self.quant_fits_only = False
+        self.deploy_target: str | None = None
+
+    def on_mount(self) -> None:
+        self.push_screen(WelcomeScreen())
+
+    # ---- guided wizard ------------------------------------------------------
+
+    def start_wizard(self) -> None:
+        self.push_screen(ModelSelectScreen())
+
+    def push_quantize_select(self) -> None:
+        self.push_screen(QuantizeSelectScreen())
+
+    def push_plan(self, candidate) -> None:  # noqa: ANN001
+        self.selected_candidate = candidate
+        self.push_screen(PlanScreen())
+
+    def push_execute(self) -> None:
+        self.push_screen(ExecuteScreen())
+
+    def go_home(self) -> None:
+        while len(self.screen_stack) > 1:
+            self.pop_screen()
+
+    # ---- settings + command mode -------------------------------------------
+
+    def push_settings(self) -> None:
+        self.push_screen(SettingsScreen())
+
+    def open_workspace(self, auto_command: str | None) -> None:
+        self.pending_command = auto_command
+        self.push_screen(WorkspaceScreen())
+
+    def resume_last_session(self) -> None:
+        latest = self.ctx.repo.latest()
+        if latest is not None:
+            latest.hardware = self.ctx.hardware
+            self.ctx.session = latest
+        self.open_workspace(None)
+
+
+def run_tui(*, session_name: str | None = None, continue_last: bool = False) -> None:
+    """Build the context and launch the TUI (blocks until the user exits)."""
+    ctx = AppContext.create(session_name=session_name, continue_last=continue_last)
+    ZeroQuantzApp(ctx=ctx).run()

zeroquantz/tui/banner.py

@@ -0,0 +1,47 @@
+"""The ``ZEROQUANTZ`` ASCII wordmark (ANSI-Shadow style), rendered responsively.
+
+Composed from per-letter glyphs so the strings can't drift out of alignment.
+Three widths: single-line (wide terminals), stacked ZERO / QUANT (typical), and a
+plain bold fallback (narrow).
+"""
+
+from __future__ import annotations
+
+_GLYPHS: dict[str, list[str]] = {
+    "Z": ["███████╗", "╚══███╔╝", "  ███╔╝ ", " ███╔╝  ", "███████╗", "╚══════╝"],
+    "E": ["███████╗", "██╔════╝", "█████╗  ", "██╔══╝  ", "███████╗", "╚══════╝"],
+    "R": ["██████╗ ", "██╔══██╗", "██████╔╝", "██╔══██╗", "██║  ██║", "╚═╝  ╚═╝"],
+    "O": [" ██████╗ ", "██╔═══██╗", "██║   ██║", "██║   ██║", "╚██████╔╝", " ╚═════╝ "],
+    "Q": [" ██████╗ ", "██╔═══██╗", "██║   ██║", "██║▄▄ ██║", "╚██████╔╝", " ╚══▀▀═╝ "],
+    "U": ["██╗   ██╗", "██║   ██║", "██║   ██║", "██║   ██║", "╚██████╔╝", " ╚═════╝ "],
+    "A": [" █████╗ ", "██╔══██╗", "███████║", "██╔══██║", "██║  ██║", "╚═╝  ╚═╝"],
+    "N": ["███╗   ██╗", "████╗  ██║", "██╔██╗ ██║", "██║╚██╗██║", "██║ ╚████║", "╚═╝  ╚═══╝"],
+    "T": ["████████╗", "╚══██╔══╝", "   ██║   ", "   ██║   ", "   ██║   ", "   ╚═╝   "],
+}
+
+
+def _compose(word: str, sep: str = " ") -> list[str]:
+    rows = ["" for _ in range(6)]
+    for ch in word:
+        glyph = _GLYPHS[ch]
+        width = max(len(r) for r in glyph)
+        for i in range(6):
+            rows[i] += glyph[i].ljust(width) + sep
+    return [r.rstrip() for r in rows]
+
+
+def banner_lines(available_width: int) -> list[str]:
+    """Return the widest banner variant that fits ``available_width``."""
+    single = _compose("ZEROQUANTZ")
+    if available_width >= max(len(r) for r in single) + 2:
+        return single
+    zero = _compose("ZERO")
+    quantz = _compose("QUANTZ")
+    stacked = zero + [""] + quantz
+    if available_width >= max(len(r) for r in stacked) + 2:
+        return stacked
+    return ["ZEROQUANTZ"]
+
+
+def banner_text(available_width: int) -> str:
+    return "\n".join(banner_lines(available_width))

zeroquantz/tui/screens/__init__.py

@@ -0,0 +1,25 @@
+"""ZeroQuantz TUI screens."""
+
+from __future__ import annotations
+
+from zeroquantz.tui.screens.confirm import ConfirmScreen
+from zeroquantz.tui.screens.execute import ExecuteScreen
+from zeroquantz.tui.screens.model_select import ModelSelectScreen
+from zeroquantz.tui.screens.plan import PlanScreen
+from zeroquantz.tui.screens.quantize_select import QuantizeSelectScreen
+from zeroquantz.tui.screens.settings import SettingsScreen
+from zeroquantz.tui.screens.token import TokenScreen
+from zeroquantz.tui.screens.welcome import WelcomeScreen
+from zeroquantz.tui.screens.workspace import WorkspaceScreen
+
+__all__ = [
+    "ConfirmScreen",
+    "ExecuteScreen",
+    "ModelSelectScreen",
+    "PlanScreen",
+    "QuantizeSelectScreen",
+    "SettingsScreen",
+    "TokenScreen",
+    "WelcomeScreen",
+    "WorkspaceScreen",
+]

zeroquantz/tui/screens/confirm.py

@@ -0,0 +1,41 @@
+"""A small yes/no confirmation modal (used before destructive actions)."""
+
+from __future__ import annotations
+
+from rich.text import Text
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import Vertical
+from textual.screen import ModalScreen
+from textual.widgets import Static
+
+
+class ConfirmScreen(ModalScreen[bool]):
+    BINDINGS = [
+        Binding("y", "confirm", "Yes"),
+        Binding("enter", "confirm", "Yes", show=False),
+        Binding("n", "cancel", "No"),
+        Binding("escape", "cancel", "No", show=False),
+    ]
+
+    def __init__(self, message: str, *, detail: str | None = None) -> None:
+        super().__init__()
+        self._message = message
+        self._detail = detail
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="confirm-dialog"):
+            yield Static(Text(self._message, style="bold white"), id="confirm-message")
+            if self._detail:
+                yield Static(Text(self._detail, style="#8b949e"), classes="dialog-sub")
+            yield Static(
+                Text.assemble(("y", "#f85149"), (" confirm    ", "#8b949e"),
+                              ("n", "#3fb950"), (" cancel", "#8b949e")),
+                classes="page-hint",
+            )
+
+    def action_confirm(self) -> None:
+        self.dismiss(True)
+
+    def action_cancel(self) -> None:
+        self.dismiss(False)

zeroquantz/tui/screens/execute.py

@@ -0,0 +1,194 @@
+"""Execute screen: run the decided plan as stages — download → quantize → deploy.
+
+Each stage streams progress/logs and shows a status glyph. Failures are surfaced
+(with the actionable message) but don't abort the pipeline — deploy assets are
+still generated so the user has runnable commands for their server.
+"""
+
+from __future__ import annotations
+
+import time
+
+from rich.text import Text
+from textual import work
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import Vertical
+from textual.screen import Screen
+from textual.widgets import Footer, ProgressBar, RichLog, Static
+
+from zeroquantz.agent.intents import Intent, IntentKind
+from zeroquantz.quantization.catalog import get_format
+from zeroquantz.utils import units
+
+_GLYPH = {"pending": ("○", "#6e7681"), "running": ("●", "#d29922"),
+          "done": ("✓", "#3fb950"), "failed": ("✗", "#f85149"), "manual": ("⤼", "#d29922")}
+
+
+class ExecuteScreen(Screen):
+    BINDINGS = [Binding("escape", "back", "Back"), Binding("h", "home", "Home")]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._stages = [
+            {"name": "Download weights", "status": "pending"},
+            {"name": "Quantize", "status": "pending"},
+            {"name": "Deploy", "status": "pending"},
+        ]
+        self._total = 0
+        self._start = 0.0
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="exec-page"):
+            yield Static("Executing plan", classes="page-title")
+            yield Static("", id="exec-stages")
+            yield ProgressBar(id="exec-progress", total=100, show_eta=False)
+            yield RichLog(id="exec-log", wrap=True, markup=False, highlight=False)
+        yield Footer()
+
+    def on_mount(self) -> None:
+        self._render_stages()
+        self._run()
+
+    # ---- stage rendering ----------------------------------------------------
+
+    def _render_stages(self) -> None:
+        text = Text()
+        for i, st in enumerate(self._stages):
+            glyph, color = _GLYPH[st["status"]]
+            text.append(f"  {glyph} ", style=color)
+            text.append(st["name"], style="bold white" if st["status"] == "running" else "white")
+            if i < len(self._stages) - 1:
+                text.append("\n")
+        self.query_one("#exec-stages", Static).update(text)
+
+    def _set_stage(self, index: int, status: str) -> None:
+        self._stages[index]["status"] = status
+        self._render_stages()
+
+    def _log(self, message: str, style: str = "#c9d1d9") -> None:
+        self.query_one("#exec-log", RichLog).write(Text(message, style=style))
+
+    def _progress(self, done: int, total: int) -> None:
+        pb = self.query_one("#exec-progress", ProgressBar)
+        pb.update(total=max(total, 1), progress=done)
+
+    def _indeterminate(self) -> None:
+        self.query_one("#exec-progress", ProgressBar).update(total=None)
+
+    # ---- pipeline -----------------------------------------------------------
+
+    @work(thread=True, exclusive=True)
+    def _run(self) -> None:
+        app = self.app
+        ctx = app.ctx
+        candidate = app.selected_candidate
+        fmt = get_format(candidate.format_id)
+        model_id = ctx.session.model_id or (ctx.session.model_profile.model_id if ctx.session.model_profile else "")
+
+        # ---- stage 1: download ----
+        app.call_from_thread(self._set_stage, 0, "running")
+        try:
+            local = ctx.session.model_local_path
+            from pathlib import Path as _P
+
+            if local and _P(local).exists():
+                app.call_from_thread(self._log, f"Weights already present: {local}", "#6e7681")
+            else:
+                from zeroquantz.models import downloader as dl
+                from zeroquantz.models.hf_auth import current_token
+
+                token = current_token()
+                files, total = dl.plan_download(model_id, token=token)
+                self._total = total
+                self._start = time.perf_counter()
+                app.call_from_thread(self._log, f"Downloading {len(files)} files · {units.humanize_bytes(total)}")
+                last = [0.0]
+
+                def on_prog(done: int) -> None:
+                    now = time.perf_counter()
+                    if done >= self._total or now - last[0] >= 0.1:
+                        last[0] = now
+                        app.call_from_thread(self._progress, done, self._total)
+
+                local = dl.download_model(model_id, files, token=token, on_progress=on_prog)
+                ctx.session.model_local_path = local
+                app.call_from_thread(self._log, f"Downloaded to {local}", "#3fb950")
+            app.call_from_thread(self._set_stage, 0, "done")
+        except Exception as exc:
+            app.call_from_thread(self._log, f"Download failed: {exc}", "#f85149")
+            app.call_from_thread(self._set_stage, 0, "failed")
+
+        # ---- stage 2: quantize ----
+        app.call_from_thread(self._set_stage, 1, "running")
+        app.call_from_thread(self._indeterminate)
+        quant_out = None
+        try:
+            def qprog(stage: str, frac: float) -> None:
+                app.call_from_thread(self._log, f"  {stage} … {frac * 100:.0f}%", "#6e7681")
+
+            result = app.dispatcher.dispatch(
+                Intent(IntentKind.QUANTIZE, {"format_id": candidate.format_id}), ctx, qprog
+            )
+            if result.success:
+                quant_out = getattr(result.payload, "output_dir", None)
+                app.call_from_thread(self._log, result.message, "#3fb950")
+                app.call_from_thread(self._set_stage, 1, "done")
+            else:
+                app.call_from_thread(self._log, result.message, "#d29922")
+                app.call_from_thread(self._set_stage, 1, "manual")
+        except Exception as exc:
+            app.call_from_thread(self._log, f"Quantize error: {exc}", "#f85149")
+            app.call_from_thread(self._set_stage, 1, "failed")
+
+        # ---- stage 3: deploy ----
+        app.call_from_thread(self._set_stage, 2, "running")
+        try:
+            from pathlib import Path
+
+            from zeroquantz.deploy import (
+                can_provision,
+                generate_assets,
+                get_target,
+                is_runtime_available,
+                launch,
+            )
+
+            target = get_target(app.deploy_target)
+            deploy_model = quant_out or ctx.session.model_local_path or model_id
+            name = model_id.split("/")[-1] or "model"
+            out = Path.cwd() / f"zeroquantz-{name}-{fmt.id}" / "deploy"
+            bundle = generate_assets(deploy_model, fmt, target, out)
+            app.call_from_thread(self._log, f"Deploy assets → {bundle.directory}", "#3fb950")
+            app.call_from_thread(self._log, f"  serve: {bundle.serve_command}", "#3fb950")
+            if bundle.endpoint:
+                app.call_from_thread(self._log, f"  endpoint (once running): {bundle.endpoint}", "#8b949e")
+            if is_runtime_available(target):
+                app.call_from_thread(self._log, f"Launching {target.label} locally…", "#d29922")
+                try:
+                    launch(target, bundle.serve_command, cuda_version=ctx.hardware.cuda_version)
+                    app.call_from_thread(self._log, f"Started — {bundle.endpoint or 'see logs'}", "#3fb950")
+                except Exception as exc:
+                    app.call_from_thread(self._log, f"Launch skipped: {exc}", "#6e7681")
+            elif can_provision(target):
+                app.call_from_thread(
+                    self._log,
+                    f"{target.label} runs in its own isolated env (not built yet). Provision + serve:",
+                    "#8b949e",
+                )
+                app.call_from_thread(self._log, f"  zeroquantz serve {deploy_model} -r {target.name}", "#58a6ff")
+                app.call_from_thread(self._log, "  (or copy the assets to a GPU server)", "#6e7681")
+            else:
+                app.call_from_thread(self._log, f"{target.label} not installed here — copy the assets to your GPU server and run.", "#8b949e")
+            app.call_from_thread(self._set_stage, 2, "done")
+        except Exception as exc:
+            app.call_from_thread(self._log, f"Deploy failed: {exc}", "#f85149")
+            app.call_from_thread(self._set_stage, 2, "failed")
+
+        app.call_from_thread(self._log, "\nPlan executed. Esc to go back, h for home.", "bold #58a6ff")
+
+    def action_back(self) -> None:
+        self.app.pop_screen()
+
+    def action_home(self) -> None:
+        self.app.go_home()