mirror of https://github.com/razor-ai/soup.git
244 lines
7.4 KiB
Python
244 lines
7.4 KiB
Python
"""Quant-Lobotomy Checker (v0.26.0 Part D).
|
|
|
|
Orchestrates a before/after eval comparison: run the same suite against two
|
|
model backends (safetensors / gguf / awq / gptq / mlx) and surface the delta.
|
|
Focused on the 'did quantization eat accuracy?' question.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
|
|
|
|
from soup_cli.utils.paths import is_under_cwd as _is_under_cwd
|
|
|
|
if TYPE_CHECKING:
|
|
from rich.table import Table
|
|
|
|
|
|
def classify_delta(delta: float, *, minor: float = 0.02, major: float = 0.05) -> str:
|
|
"""Return one of OK / MINOR / MAJOR based on absolute drop in score."""
|
|
if delta >= 0:
|
|
return "OK"
|
|
drop = -delta
|
|
if drop < minor:
|
|
return "OK"
|
|
if drop < major:
|
|
return "MINOR"
|
|
return "MAJOR"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class QuantCheckRow:
|
|
task: str
|
|
before: float
|
|
after: float
|
|
delta: float
|
|
verdict: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class QuantCheckResult:
|
|
rows: list[QuantCheckRow]
|
|
|
|
def to_json(self) -> str:
|
|
return json.dumps(
|
|
{"rows": [asdict(r) for r in self.rows]},
|
|
indent=2,
|
|
)
|
|
|
|
|
|
def _score_tasks(tasks_file: str, generate_fn: Callable[[str], str]) -> float:
|
|
"""Load a JSONL task file and return the average score."""
|
|
from soup_cli.eval.custom import load_eval_tasks, score_task
|
|
|
|
tasks = load_eval_tasks(tasks_file)
|
|
if not tasks:
|
|
return 0.0
|
|
total = 0.0
|
|
for task in tasks:
|
|
output = generate_fn(task.prompt)
|
|
total += float(score_task(task, output).score)
|
|
return total / len(tasks)
|
|
|
|
|
|
def run_quant_check(
|
|
*,
|
|
before_gen: Callable[[str], str],
|
|
after_gen: Callable[[str], str],
|
|
tasks_file: str,
|
|
task_name: str = "default",
|
|
) -> QuantCheckResult:
|
|
"""Run ``tasks_file`` through both models and return the row."""
|
|
before = _score_tasks(tasks_file, before_gen)
|
|
after = _score_tasks(tasks_file, after_gen)
|
|
delta = after - before
|
|
verdict = classify_delta(delta)
|
|
return QuantCheckResult(rows=[QuantCheckRow(
|
|
task=task_name, before=before, after=after,
|
|
delta=delta, verdict=verdict,
|
|
)])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rich rendering helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def render_table(result: QuantCheckResult) -> "Table":
|
|
"""Render a QuantCheckResult as a Rich Table (for the CLI)."""
|
|
from rich.table import Table
|
|
|
|
table = Table(title="Quant check")
|
|
table.add_column("Task", style="cyan")
|
|
table.add_column("Before", justify="right")
|
|
table.add_column("After", justify="right")
|
|
table.add_column("Delta", justify="right")
|
|
table.add_column("Verdict")
|
|
for row in result.rows:
|
|
verdict = row.verdict
|
|
colored = {
|
|
"OK": f"[green]{verdict}[/]",
|
|
"MINOR": f"[yellow]{verdict}[/]",
|
|
"MAJOR": f"[red]{verdict}[/]",
|
|
}.get(verdict, verdict)
|
|
table.add_row(
|
|
row.task,
|
|
f"{row.before:.3f}",
|
|
f"{row.after:.3f}",
|
|
f"{row.delta:+.3f}",
|
|
colored,
|
|
)
|
|
return table
|
|
|
|
|
|
def render_markdown(result: QuantCheckResult) -> str:
|
|
lines = [
|
|
"| Task | Before | After | Delta | Verdict |",
|
|
"|------|--------|-------|-------|---------|",
|
|
]
|
|
for row in result.rows:
|
|
lines.append(
|
|
f"| {row.task} | {row.before:.3f} | {row.after:.3f} | "
|
|
f"{row.delta:+.3f} | {row.verdict} |"
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
|
|
FORMAT_RENDERERS: dict[str, Callable[[QuantCheckResult], Union["Table", str]]] = {
|
|
"table": render_table,
|
|
"json": lambda r: r.to_json(),
|
|
"markdown": render_markdown,
|
|
}
|
|
|
|
|
|
def render(
|
|
result: QuantCheckResult, *, fmt: Literal["table", "json", "markdown"],
|
|
) -> Union["Table", str]:
|
|
renderer = FORMAT_RENDERERS.get(fmt)
|
|
if renderer is None:
|
|
raise ValueError(f"unknown format '{fmt}'. Use table | json | markdown")
|
|
return renderer(result)
|
|
|
|
|
|
def ensure_format(fmt: str) -> None:
|
|
"""Validate format string, raise ValueError on unknown."""
|
|
if fmt not in FORMAT_RENDERERS:
|
|
raise ValueError(f"unknown format '{fmt}'. Use table | json | markdown")
|
|
|
|
|
|
def make_model_generator(
|
|
model_path: str,
|
|
*,
|
|
max_new_tokens: int = 256,
|
|
temperature: float = 0.0,
|
|
) -> Callable[[str], str]:
|
|
"""Return a ``generate_fn(prompt) -> str`` backed by a transformers model.
|
|
|
|
Lazy-loaded so the CLI stays cold-start fast. The model is loaded once
|
|
and reused across calls. ``temperature=0`` enables greedy decoding for
|
|
reproducible eval scores.
|
|
"""
|
|
if max_new_tokens < 1 or max_new_tokens > 16384:
|
|
raise ValueError("max_new_tokens must be in [1, 16384]")
|
|
|
|
from transformers import (
|
|
AutoModelForCausalLM,
|
|
AutoTokenizer,
|
|
)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=False)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_path, trust_remote_code=False
|
|
)
|
|
model.eval()
|
|
|
|
def _generate(prompt: str) -> str:
|
|
if not prompt:
|
|
return ""
|
|
inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
|
|
do_sample = temperature > 0.0
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=do_sample,
|
|
temperature=max(temperature, 1e-5),
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
# Strip the prompt prefix from the decoded text.
|
|
new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
|
|
return tokenizer.decode(new_tokens, skip_special_tokens=True)
|
|
|
|
return _generate
|
|
|
|
|
|
def stub_generator(label: str) -> Callable[[str], str]:
|
|
"""Return a deterministic stub generator so the CLI has something runnable.
|
|
|
|
Real model loading is wired post-v0.26.0 — once then, pass in real
|
|
``before_gen`` / ``after_gen`` callables instead of this stub.
|
|
"""
|
|
def _stub(prompt: str) -> str: # noqa: ARG001
|
|
return f"[stub:{label}]"
|
|
return _stub
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Path containment (realpath + commonpath)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def is_under_cwd(path: Path) -> bool:
|
|
"""Backwards-compat alias — use :func:`soup_cli.utils.paths.is_under_cwd`."""
|
|
return _is_under_cwd(path)
|
|
|
|
|
|
def resolve_model_ref(
|
|
ref: str, *, kinds: Optional[tuple[str, ...]] = None,
|
|
) -> Optional[str]:
|
|
"""Resolve ``registry://<id>`` to an artifact path, or pass through a file.
|
|
|
|
When ``kinds`` is provided, only artifacts whose ``kind`` field matches
|
|
are returned. This prevents picking the wrong artifact when an entry
|
|
has multiple attachments (e.g. adapter + GGUF + safetensors).
|
|
"""
|
|
if ref.startswith("registry://"):
|
|
from soup_cli.registry.store import RegistryStore
|
|
|
|
with RegistryStore() as store:
|
|
eid = store.resolve(ref)
|
|
if eid is None:
|
|
return None
|
|
arts = store.get_artifacts(eid)
|
|
if not arts:
|
|
return None
|
|
if kinds:
|
|
filtered = [a for a in arts if a.get("kind") in kinds]
|
|
if not filtered:
|
|
return None
|
|
return filtered[0]["path"]
|
|
return arts[0]["path"]
|
|
return ref
|