quantdiff 0.1.0rc1__py3-none-any.whl

quantdiff/preflight.py

@@ -0,0 +1,365 @@
+"""Pre-flight checks that catch setup mistakes before a long comparison run.
+
+Each check is isolated: a server error inside one check becomes a "skip" finding for
+that check instead of aborting the others. Checks that do not apply to the given
+options (no reference, no Hugging Face repo, probe disabled) produce no finding.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import re
+from collections.abc import Callable, Mapping, Sequence
+from typing import Final
+
+from quantdiff._http import get_json
+from quantdiff.backends.base import Backend
+from quantdiff.errors import BackendError, SpecError
+from quantdiff.types import BackendKind, Message, PreflightFinding, ServerInfo
+
+__all__ = [
+    "DEFAULT_CONTEXT_PROBE_TOKENS",
+    "HF_BASE_URL",
+    "MIN_CONTEXT_PROBE_TOKENS",
+    "run_preflight",
+]
+
+logger = logging.getLogger(__name__)
+
+HF_BASE_URL: Final = "https://huggingface.co"
+DEFAULT_CONTEXT_PROBE_TOKENS: Final = 6000
+MIN_CONTEXT_PROBE_TOKENS: Final = 1000
+"""Smallest long-probe size worth running; 0 disables the probe instead."""
+
+_CONTROL_PROBE_TOKENS: Final = 200
+_WORDS_PER_TOKEN: Final = 0.75
+_PROBE_ANSWER_TOKENS: Final = 16
+_HF_TIMEOUT_SECONDS: Final = 15.0
+_HF_MAX_BYTES: Final = 2 * 1024 * 1024
+_HF_REPO_PATTERN: Final = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*")
+
+_CONTEXT_FIXES: Final[Mapping[BackendKind, str]] = {
+    "ollama": (
+        "Set OLLAMA_CONTEXT_LENGTH={size} on the Ollama server, or add "
+        "PARAMETER num_ctx {size} to a Modelfile."
+    ),
+    "llamacpp": "Restart llama-server with -c {size} (--ctx-size).",
+    "openai": (
+        "Raise the context length to at least {size} tokens in the server settings "
+        "(LM Studio: Context Length; vLLM: --max-model-len)."
+    ),
+}
+_TEMPLATE_FIX: Final = (
+    "Re-download the model, or pass --chat-template-file with the upstream template."
+)
+_LOGPROBS_FIX: Final = "Serve the model with llama-server or another backend that returns logprobs."
+_TOKENIZER_FIX: Final = "Compare quantizations of the same base model as the reference."
+
+_TOKENIZER_PROBES: Final = (
+    "The quick brown fox jumps over the lazy dog.",
+    "  leading spaces,\ttabs\nand newlines\n\n",
+    "def add(a: int, b: int) -> int:\n    return a + b  # sum\n",
+    '{"id": 42, "tags": ["x", "y"], "ok": true}',
+    "Grüße aus München, café crème brûlée",
+    "你好世界。量化模型可以在普通电脑上运行",
+    "नमस्ते दुनिया",
+    "emoji \U0001f642\U0001f680 and numbers 3.14159 -1024 1e-9",
+)
+
+_SYLLABLES: Final = (
+    "ka", "lo", "mi", "ru", "te", "va", "zo", "ne",
+    "pi", "su", "do", "fa", "gi", "ho", "ju", "be",
+)  # fmt: skip
+_SUBJECTS: Final = (
+    "The harbor master",
+    "A patient gardener",
+    "The night librarian",
+    "An old cartographer",
+    "The village baker",
+    "A young engineer",
+    "The ferry captain",
+    "A careful auditor",
+    "The museum guide",
+    "A retired teacher",
+    "The orchard keeper",
+    "A traveling violinist",
+    "The station clerk",
+)
+_VERBS: Final = (
+    "quietly repaired",
+    "carefully counted",
+    "slowly painted",
+    "proudly displayed",
+    "briefly studied",
+    "neatly folded",
+    "gently polished",
+    "patiently sorted",
+    "happily described",
+    "firmly secured",
+    "openly admired",
+)
+_OBJECTS: Final = (
+    "a wooden bench",
+    "the brass lanterns",
+    "several faded maps",
+    "a basket of pears",
+    "the winter blankets",
+    "an iron gate",
+    "a stack of letters",
+)
+_PLACES: Final = (
+    "near the river before noon",
+    "beside the old market square",
+    "in the garden after the rain",
+    "at the edge of the quiet town",
+    "under the tall chestnut trees",
+)
+_SENTENCES_PER_PARAGRAPH: Final = 8
+
+
+def run_preflight(
+    backend: Backend,
+    *,
+    reference: Backend | None = None,
+    hf_repo: str | None = None,
+    offline: bool = False,
+    context_probe_tokens: int = DEFAULT_CONTEXT_PROBE_TOKENS,
+    hf_base_url: str = HF_BASE_URL,
+) -> tuple[PreflightFinding, ...]:
+    """Run the pre-flight checks for `backend` and return their findings in order.
+
+    `context_probe_tokens` sets the size of the long-prompt truncation probe; 0 turns
+    the context check off. `hf_repo` (such as "Qwen/Qwen2.5-7B-Instruct") enables the
+    chat template comparison. `reference` enables the tokenizer comparison.
+    """
+    if hf_repo is not None and not _HF_REPO_PATTERN.fullmatch(hf_repo):
+        raise SpecError(f"invalid Hugging Face repo {hf_repo!r}; expected owner/name")
+    if context_probe_tokens != 0 and context_probe_tokens < MIN_CONTEXT_PROBE_TOKENS:
+        raise SpecError(
+            f"context probe must be 0 (off) or at least {MIN_CONTEXT_PROBE_TOKENS} tokens"
+        )
+
+    findings = _isolated("logprobs", lambda: (_check_logprobs(backend),))
+    if context_probe_tokens:
+        findings += _isolated("context", lambda: _check_context(backend, context_probe_tokens))
+    if hf_repo is not None:
+        findings += _isolated(
+            "template",
+            lambda: (_check_template(backend, hf_repo, offline=offline, base_url=hf_base_url),),
+        )
+    if reference is not None:
+        findings += _isolated("tokenizer", lambda: (_check_tokenizer(backend, reference),))
+    for finding in findings:
+        logger.debug("preflight %s %s: %s", finding.check, finding.severity, finding.message)
+    return tuple(findings)
+
+
+def _isolated(check: str, run: Callable[[], Sequence[PreflightFinding]]) -> list[PreflightFinding]:
+    try:
+        return list(run())
+    except BackendError as exc:
+        return [PreflightFinding(check, "skip", f"check could not run: {exc}")]
+
+
+# logprobs --------------------------------------------------------------------------------
+
+
+def _check_logprobs(backend: Backend) -> PreflightFinding:
+    if backend.info().supports_logprobs:
+        return PreflightFinding("logprobs", "ok", "logprobs available")
+    return PreflightFinding(
+        "logprobs", "warn", "logit metrics unavailable; task metrics only", _LOGPROBS_FIX
+    )
+
+
+# context ---------------------------------------------------------------------------------
+
+
+def _check_context(backend: Backend, probe_tokens: int) -> tuple[PreflightFinding, ...]:
+    info = backend.info()
+    fix = _CONTEXT_FIXES[info.backend].format(size=_suggested_context(probe_tokens))
+    findings = []
+    if info.context_length is not None and info.context_length < probe_tokens:
+        findings.append(
+            PreflightFinding(
+                "context",
+                "warn",
+                f"context window is {info.context_length} tokens, below the "
+                f"{probe_tokens}-token probe; long prompts will not fit",
+                fix,
+            )
+        )
+    findings.append(_needle_probe(backend, probe_tokens, fix))
+    return tuple(findings)
+
+
+def _needle_probe(backend: Backend, probe_tokens: int, fix: str) -> PreflightFinding:
+    """Check that text at the very start of a long prompt still reaches the model.
+
+    Servers that silently drop the front of an over-long prompt return fluent answers
+    that are simply wrong, so the probe hides a code word at the start and asks for it
+    at the end. A short control probe first rules out models that cannot do the task.
+    """
+    try:
+        control_passed = _recalls_code_word(backend, _CONTROL_PROBE_TOKENS, "control")
+    except BackendError as exc:
+        return PreflightFinding("context", "skip", f"short probe failed: {exc}")
+    if not control_passed:
+        return PreflightFinding(
+            "context",
+            "skip",
+            "model could not answer the probe at short length; context check inconclusive",
+        )
+    try:
+        long_passed = _recalls_code_word(backend, probe_tokens, "long")
+    except BackendError as exc:
+        return PreflightFinding("context", "warn", f"{probe_tokens}-token probe failed: {exc}", fix)
+    if not long_passed:
+        return PreflightFinding(
+            "context", "fail", "front of long prompts is being dropped (silent truncation)", fix
+        )
+    return PreflightFinding(
+        "context", "ok", f"model recalled the start of a {probe_tokens}-token prompt"
+    )
+
+
+def _recalls_code_word(backend: Backend, probe_tokens: int, salt: str) -> bool:
+    code_word = _code_word(salt)
+    prompt = (
+        f"Remember this code word: {code_word}.\n\n"
+        f"{_filler(int(probe_tokens * _WORDS_PER_TOKEN))}\n\n"
+        "What was the code word given at the very start of this message? "
+        "Reply with the code word only."
+    )
+    result = backend.chat([Message(role="user", content=prompt)], max_tokens=_PROBE_ANSWER_TOKENS)
+    return _alphanumeric(code_word) in _alphanumeric(result.text)
+
+
+def _code_word(salt: str) -> str:
+    digest = hashlib.sha256(f"quantdiff-needle-{salt}".encode()).digest()
+    word = "".join(_SYLLABLES[byte % len(_SYLLABLES)] for byte in digest[:4])
+    number = int.from_bytes(digest[4:6], "big") % 9000 + 1000
+    return f"{word.capitalize()}-{number}"
+
+
+def _filler(word_count: int) -> str:
+    """Return about `word_count` words of varied, deterministic, meaningless prose."""
+    paragraphs: list[str] = []
+    sentences: list[str] = []
+    words = 0
+    index = 0
+    while words < word_count:
+        sentence = (
+            f"{_SUBJECTS[index % len(_SUBJECTS)]} {_VERBS[index % len(_VERBS)]} "
+            f"{_OBJECTS[index % len(_OBJECTS)]} {_PLACES[index % len(_PLACES)]}."
+        )
+        sentences.append(sentence)
+        words += len(sentence.split())
+        index += 1
+        if len(sentences) == _SENTENCES_PER_PARAGRAPH:
+            paragraphs.append(" ".join(sentences))
+            sentences = []
+    if sentences:
+        paragraphs.append(" ".join(sentences))
+    return "\n\n".join(paragraphs)
+
+
+def _alphanumeric(text: str) -> str:
+    return re.sub(r"[^a-z0-9]", "", text.lower())
+
+
+def _suggested_context(probe_tokens: int) -> int:
+    """Smallest power of two that fits the probe plus room for an answer."""
+    size = 1024
+    while size < probe_tokens + 512:
+        size *= 2
+    return size
+
+
+# chat template ---------------------------------------------------------------------------
+
+
+def _check_template(
+    backend: Backend, repo: str, *, offline: bool, base_url: str
+) -> PreflightFinding:
+    info = backend.info()
+    reason = _template_skip_reason(info, repo, offline=offline)
+    if reason is not None:
+        return PreflightFinding("template", "skip", reason)
+    upstream = _upstream_template(repo, base_url)
+    if upstream is None:
+        return PreflightFinding(
+            "template", "skip", f"upstream {repo} tokenizer_config.json has no chat_template"
+        )
+    if _squash(upstream) == _squash(info.chat_template or ""):
+        return PreflightFinding("template", "ok", f"embedded chat template matches upstream {repo}")
+    return PreflightFinding(
+        "template", "warn", f"embedded chat template differs from upstream {repo}", _TEMPLATE_FIX
+    )
+
+
+def _template_skip_reason(info: ServerInfo, repo: str, *, offline: bool) -> str | None:
+    if info.template_dialect == "go":
+        return "Ollama templates are Go templates; compare not supported"
+    if info.template_dialect != "jinja":
+        return "server does not report a Jinja chat template"
+    if offline:
+        return f"offline; upstream template for {repo} not fetched"
+    if info.chat_template is None:
+        return "server did not report its chat template"
+    return None
+
+
+def _upstream_template(repo: str, base_url: str) -> str | None:
+    url = f"{base_url.rstrip('/')}/{repo}/raw/main/tokenizer_config.json"
+    config = get_json(url, timeout=_HF_TIMEOUT_SECONDS, max_bytes=_HF_MAX_BYTES)
+    if not isinstance(config, dict):
+        raise BackendError(f"{url} did not return a JSON object")
+    template = config.get("chat_template")
+    if isinstance(template, str):
+        return template
+    if isinstance(template, list):
+        for entry in template:
+            if (
+                isinstance(entry, dict)
+                and entry.get("name") == "default"
+                and isinstance(entry.get("template"), str)
+            ):
+                return str(entry["template"])
+    return None
+
+
+def _squash(text: str) -> str:
+    return " ".join(text.split())
+
+
+# tokenizer -------------------------------------------------------------------------------
+
+
+def _check_tokenizer(backend: Backend, reference: Backend) -> PreflightFinding:
+    for probe in _TOKENIZER_PROBES:
+        candidate_ids = backend.tokenize(probe)
+        if candidate_ids is None:
+            return _tokenizer_skip(backend)
+        reference_ids = reference.tokenize(probe)
+        if reference_ids is None:
+            return _tokenizer_skip(reference)
+        if candidate_ids != reference_ids:
+            return PreflightFinding(
+                "tokenizer",
+                "fail",
+                "tokenizers differ from the reference; logit metrics are not comparable",
+                _TOKENIZER_FIX,
+            )
+    return PreflightFinding(
+        "tokenizer",
+        "ok",
+        f"tokenizer matches the reference on {len(_TOKENIZER_PROBES)} probe strings",
+    )
+
+
+def _tokenizer_skip(backend: Backend) -> PreflightFinding:
+    return PreflightFinding(
+        "tokenizer", "skip", f"{backend.spec.label} cannot tokenize; tokenizers not compared"
+    )

quantdiff/progress.py

@@ -0,0 +1,283 @@
+"""Progress displays for long runs.
+
+`make_progress` picks one of two displays for a stream:
+
+- On a terminal, one status line is redrawn in place with a bar, percentage, ETA and the
+  current step. Each finished step leaves a permanent `done` line so the scrollback reads
+  as a short log. A step is finished only when the next one starts, so a step that fails
+  never gets a `done` line, and connecting to every server counts as one step whose
+  `done` line says how many servers answered.
+- On a pipe or in CI, plain lines are printed when the step changes and at most once per
+  tenth of the run, so logs stay short and contain no carriage returns.
+
+The ETA divides the units left by the rate measured in the current step. Work units are
+weighted a priori and a scoring unit costs far less than a case, so a rate averaged over
+the whole run is badly off; the recent rate of the step at hand is a much better guide.
+
+Output is ASCII only so it renders on any Windows console code page.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Final, TextIO
+
+from quantdiff.types import ProgressEvent
+
+Clock = Callable[[], float]
+ProgressFn = Callable[[ProgressEvent], None]
+
+BAR_WIDTH: Final = 20
+REDRAW_INTERVAL_SECONDS: Final = 0.1
+ETA_MIN_UNITS: Final = 5
+ETA_MIN_SECONDS: Final = 10.0
+"""Units and seconds into a step before its ETA is shown; earlier estimates swing wildly."""
+RATE_WINDOW_EVENTS: Final = 10
+"""The measured rate weights roughly the last this many events."""
+LOG_STEPS: Final = 10
+"""Non-interactive output prints at most once per 1/LOG_STEPS of the run."""
+
+
+def make_progress(
+    stream: TextIO, *, interactive: bool | None = None, clock: Clock = time.monotonic
+) -> ProgressDisplay:
+    """Return a progress callback writing to `stream`.
+
+    `interactive` None means a live status line when `stream` is a terminal and TERM is not
+    "dumb", and plain log lines otherwise. Call `close()` on the result before printing an
+    error so a half-drawn status line does not run into it.
+    """
+    if interactive is None:
+        interactive = stream.isatty() and os.environ.get("TERM") != "dumb"
+    if interactive:
+        return StatusLine(stream, clock)
+    return StepLog(stream, clock)
+
+
+def format_duration(seconds: float) -> str:
+    """Format as 45s, 3m05s or 1h02m."""
+    whole = max(0, round(seconds))
+    if whole < 60:
+        return f"{whole}s"
+    minutes, secs = divmod(whole, 60)
+    if minutes < 60:
+        return f"{minutes}m{secs:02d}s"
+    hours, minutes = divmod(minutes, 60)
+    return f"{hours}h{minutes:02d}m"
+
+
+class RateEstimator:
+    """Units per second within one step, smoothed over the recent events.
+
+    Units and seconds are smoothed separately and divided, so a burst of events that
+    arrive together (a cache hit) does not count as an infinitely fast interval.
+    """
+
+    def __init__(self) -> None:
+        self._alpha = 2 / (RATE_WINDOW_EVENTS + 1)
+        self._started = 0.0
+        self._start_units = 0
+        self._last = 0.0
+        self._last_units = 0
+        self._units = 0.0
+        self._seconds = 0.0
+
+    def start(self, now: float, completed: int) -> None:
+        """Begin a new step whose work started at `now` with `completed` units done."""
+        self._started = self._last = now
+        self._start_units = self._last_units = completed
+        self._units = self._seconds = 0.0
+
+    def update(self, now: float, completed: int) -> None:
+        units = max(0, completed - self._last_units)
+        seconds = max(0.0, now - self._last)
+        self._units += self._alpha * (units - self._units)
+        self._seconds += self._alpha * (seconds - self._seconds)
+        self._last, self._last_units = now, completed
+
+    def remaining(self, completed: int, total: int) -> float | None:
+        """Seconds left at the recent rate, or None while too little of the step is done."""
+        if total <= 0 or completed >= total:
+            return None
+        warming_up = (
+            self._last - self._started < ETA_MIN_SECONDS
+            or self._last_units - self._start_units < ETA_MIN_UNITS
+        )
+        if warming_up or self._units <= 0 or self._seconds <= 0:
+            return None
+        return (total - completed) * self._seconds / self._units
+
+
+class ProgressDisplay(ABC):
+    """Shared step tracking: notices when the (phase, model) step changes, times it and
+    measures its rate for the ETA."""
+
+    def __init__(self, stream: TextIO, clock: Clock) -> None:
+        self._stream = stream
+        self._clock = clock
+        self._run_started: float | None = None
+        self._step: tuple[str, str] | None = None
+        self._step_started = 0.0
+        self._step_events = 0
+        self._rate = RateEstimator()
+        self._last_event: tuple[float, int] | None = None
+
+    def __call__(self, event: ProgressEvent) -> None:
+        now = self._clock()
+        if self._run_started is None:
+            self._run_started = now
+        if event.phase == "done":
+            self._finish_step(now)
+            self._finish_run(now - self._run_started)
+            return
+        step = _step_key(event)
+        new_step = step != self._step
+        if new_step:
+            self._finish_step(now)
+            self._step = step
+            self._step_started = now
+            self._step_events = 0
+            # Events report work already done, so the step's first unit began at the
+            # previous event.
+            self._rate.start(*(self._last_event or (now, event.completed)))
+        self._step_events += 1
+        self._rate.update(now, event.completed)
+        self._last_event = (now, event.completed)
+        if new_step:
+            self._start_step(event, now)
+        else:
+            self._advance(event, now)
+
+    @abstractmethod
+    def close(self) -> None:
+        """Leave the stream at the start of a clean line. Safe to call more than once."""
+
+    def _remaining(self, event: ProgressEvent) -> float | None:
+        return self._rate.remaining(event.completed, event.total)
+
+    def _finish_step(self, now: float) -> None:
+        if self._step is None:
+            return
+        self.close()
+        phase, model = self._step
+        if phase == "connect":
+            # The runner sends one connect event per server once that server has answered.
+            label = f"connect {self._step_events} server{'' if self._step_events == 1 else 's'}"
+        else:
+            label = f"{phase} {model}".rstrip()
+        self._write_line(f"  done  {label}  ({format_duration(now - self._step_started)})")
+        self._step = None
+
+    def _finish_run(self, elapsed: float) -> None:
+        self.close()
+        self._write_line(f"Finished in {format_duration(elapsed)}")
+
+    @abstractmethod
+    def _start_step(self, event: ProgressEvent, now: float) -> None:
+        """Show the first event of a new step."""
+
+    @abstractmethod
+    def _advance(self, event: ProgressEvent, now: float) -> None:
+        """Show a later event of the current step, if it is worth showing."""
+
+    def _write_line(self, text: str) -> None:
+        self._stream.write(text + "\n")
+        self._stream.flush()
+
+
+class StatusLine(ProgressDisplay):
+    """A single line redrawn in place, for terminals."""
+
+    def __init__(self, stream: TextIO, clock: Clock) -> None:
+        super().__init__(stream, clock)
+        self._drawn_width = 0
+        self._last_draw: float | None = None
+
+    def close(self) -> None:
+        if self._drawn_width:
+            # Overwrite with spaces rather than an ANSI erase: legacy Windows consoles
+            # print escape sequences literally unless virtual terminal mode is enabled.
+            self._stream.write("\r" + " " * self._drawn_width + "\r")
+            self._stream.flush()
+            self._drawn_width = 0
+
+    def _start_step(self, event: ProgressEvent, now: float) -> None:
+        self._draw(event, now)
+
+    def _advance(self, event: ProgressEvent, now: float) -> None:
+        if self._last_draw is None or now - self._last_draw >= REDRAW_INTERVAL_SECONDS:
+            self._draw(event, now)
+
+    def _draw(self, event: ProgressEvent, now: float) -> None:
+        width = max(1, shutil.get_terminal_size().columns - 1)
+        line = status_text(event, self._remaining(event))[:width]
+        padding = " " * max(0, self._drawn_width - len(line))
+        self._stream.write("\r" + line + padding)
+        self._stream.flush()
+        self._drawn_width = len(line)
+        self._last_draw = now
+
+
+class StepLog(ProgressDisplay):
+    """Plain lines for pipes and CI logs."""
+
+    def __init__(self, stream: TextIO, clock: Clock) -> None:
+        super().__init__(stream, clock)
+        self._logged_step = -1
+
+    def close(self) -> None:
+        """Nothing to clean up: every log line already ends with a newline."""
+
+    def _start_step(self, event: ProgressEvent, now: float) -> None:
+        self._log(event)
+
+    def _advance(self, event: ProgressEvent, now: float) -> None:
+        if _log_step(event) > self._logged_step:
+            self._log(event)
+
+    def _log(self, event: ProgressEvent) -> None:
+        self._logged_step = _log_step(event)
+        parts = [f"[{_percent(event):>4}]", event.phase, event.model, event.detail]
+        remaining = self._remaining(event)
+        if remaining is not None:
+            parts.append(f"(ETA {format_duration(remaining)})")
+        self._write_line("  ".join(part for part in parts if part))
+
+
+def status_text(event: ProgressEvent, remaining: float | None) -> str:
+    """The interactive status line, before truncation to the terminal width."""
+    fraction = _fraction(event)
+    filled = 0 if fraction is None else round(fraction * BAR_WIDTH)
+    parts = [f"[{'#' * filled}{'.' * (BAR_WIDTH - filled)}]", f"{_percent(event):>4}"]
+    if remaining is not None:
+        parts.append(f"ETA {format_duration(remaining)}")
+    parts.extend((event.model, f"{event.phase} {event.detail}".rstrip()))
+    return "  ".join(part for part in parts if part)
+
+
+def _step_key(event: ProgressEvent) -> tuple[str, str]:
+    # The runner connects to every server before any work; a done line per server would
+    # time the wrong one, since each event arrives after its connection attempt.
+    if event.phase == "connect":
+        return ("connect", "")
+    return (event.phase, event.model)
+
+
+def _fraction(event: ProgressEvent) -> float | None:
+    if event.total <= 0:
+        return None
+    return min(1.0, max(0.0, event.completed / event.total))
+
+
+def _percent(event: ProgressEvent) -> str:
+    fraction = _fraction(event)
+    return "--%" if fraction is None else f"{int(fraction * 100)}%"
+
+
+def _log_step(event: ProgressEvent) -> int:
+    fraction = _fraction(event)
+    return 0 if fraction is None else int(fraction * LOG_STEPS)