logit-classifier 0.1.0__py3-none-any.whl

logit_classifier/calibrate.py

@@ -0,0 +1,141 @@
+"""Batch calibration: a running per-label prior learned from served traffic.
+
+The label prior is token bias, the model's standing preference for the token
+"A" over "B". It is estimated as a running mean of the uncalibrated label
+distribution and subtracted in log space before the softmax. No labelled data
+and no extra forward passes are required.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import math
+import threading
+from pathlib import Path
+from uuid import uuid4
+
+import numpy as np
+
+from .config import PRIOR_MIN_OBSERVATIONS, PRIOR_SMOOTHING
+
+
+def _stored_mean(key: str, vector: object) -> list[float] | None:
+    """Return the mean of one `kind:width` bucket, or None when the stored value is unusable."""
+    kind, separator, width = key.partition(":")
+
+    if not kind or separator != ":" or not (width.isascii() and width.isdigit()):
+        return None
+    if not isinstance(vector, list) or len(vector) != int(width):
+        return None
+    if not all(isinstance(value, int | float) and not isinstance(value, bool) for value in vector):
+        return None
+    values = [float(value) for value in vector]
+    if not all(math.isfinite(value) and value >= 0.0 for value in values):
+        return None
+    return values
+
+
+class PriorStore:
+    def __init__(
+        self, path: Path | None, fingerprint: str, min_observations: int = PRIOR_MIN_OBSERVATIONS
+    ) -> None:
+        self.path = path
+        self.fingerprint = fingerprint
+        self.min_observations = min_observations
+        self._lock = threading.Lock()
+        self._means: dict[str, list[float]] = {}
+        self._counts: dict[str, int] = {}
+        self._load()
+
+    @staticmethod
+    def _bucket(kind: str, label_count: int) -> str:
+        return f"{kind}:{label_count}"
+
+    def _load(self) -> None:
+        means: dict[str, list[float]] = {}
+        counts: dict[str, int] = {}
+
+        if self.path is None or not self.path.exists():
+            return
+        try:
+            payload = json.loads(self.path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            return
+        if not isinstance(payload, dict):
+            return
+        # A prior measured under a different model or prompt describes a different
+        # distribution, so a mismatch starts over rather than blending the two.
+        if payload.get("fingerprint") != self.fingerprint:
+            return
+        stored_means = payload.get("means")
+        stored_counts = payload.get("counts")
+        if not isinstance(stored_means, dict) or not isinstance(stored_counts, dict):
+            return
+
+        # The path is user-visible state, so a hand-edited bucket is a reachable input.
+        # One that does not match its own key starts over the way a mismatch does, rather
+        # than reaching observe and raising on every later request.
+        for key, vector in stored_means.items():
+            mean = _stored_mean(key, vector) if isinstance(key, str) else None
+            count = stored_counts.get(key)
+            if mean is None or not isinstance(count, int) or isinstance(count, bool) or count < 0:
+                continue
+            means[key] = mean
+            counts[key] = count
+
+        self._means = means
+        self._counts = counts
+
+    def save(self) -> None:
+        # A path with no final component, `.` or a drive root, has no staged sibling to
+        # name, and with_name raises ValueError, which the OSError suppression misses.
+        if self.path is None or not self.path.name:
+            return
+        payload = {"fingerprint": self.fingerprint, "means": self._means, "counts": self._counts}
+
+        # Two processes can hold a store over one path, and a truncating write hands a
+        # reader a partial file that loads as an empty prior. A rename is atomic.
+        with self._lock:
+            staged = self.path.with_name(f"{self.path.name}.{uuid4().hex}.tmp")
+            try:
+                staged.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+                staged.replace(self.path)
+            except OSError:
+                # A rename that keeps failing would otherwise leave one staged file per
+                # request, since the uuid name is never reused.
+                with contextlib.suppress(OSError):
+                    staged.unlink(missing_ok=True)
+
+    def observe(self, kind: str, probabilities: np.ndarray) -> None:
+        """Fold one uncalibrated distribution into the running mean."""
+        bucket = self._bucket(kind, len(probabilities))
+
+        with self._lock:
+            count = self._counts.get(bucket, 0)
+            mean = np.asarray(self._means.get(bucket, [0.0] * len(probabilities)), dtype=np.float64)
+            self._means[bucket] = ((count * mean + probabilities) / (count + 1)).tolist()
+            self._counts[bucket] = count + 1
+
+    def log_prior(self, kind: str, label_count: int) -> np.ndarray | None:
+        """Return the log-prior to subtract, or None while the estimate is thin."""
+        bucket = self._bucket(kind, label_count)
+
+        with self._lock:
+            count = self._counts.get(bucket, 0)
+            mean = self._means.get(bucket)
+        if mean is None or count < self.min_observations or len(mean) != label_count:
+            return None
+        prior = np.asarray(mean, dtype=np.float64)
+        if not np.all(prior >= 0) or prior.sum() <= 0:
+            return None
+        prior = prior / prior.sum()
+        prior = (1.0 - PRIOR_SMOOTHING) * prior + PRIOR_SMOOTHING / label_count
+        # Centring keeps the subtraction from shifting overall scale, which only
+        # the temperature should control.
+        log_prior = np.log(prior)
+        return log_prior - log_prior.mean()
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            return dict(self._counts)

logit_classifier/classifier.py

@@ -0,0 +1,329 @@
+"""Pipeline coordinator: request in, Jev-shaped answers out."""
+
+from __future__ import annotations
+
+import hashlib
+import random
+from dataclasses import dataclass, field, replace
+
+import numpy as np
+
+from .backends.base import Backend, BackendContractError, BranchLogits
+from .calibrate import PriorStore
+from .config import ANSWER_PREFILL, Config, fitted_temperature
+from .deps import MissingDependencyError
+from .prompt import (
+    PROMPT_VERSION,
+    SYSTEM_PROMPT,
+    Branch,
+    branch_content,
+    build_branches,
+    prefix_content,
+)
+from .schema import (
+    Answer,
+    ChoiceAnswer,
+    ChoiceQuestion,
+    NoulAnswer,
+    NoulQuestion,
+    Question,
+    ScoreAnswer,
+    ScoreQuestion,
+    SystemOneRequest,
+    SystemOneResponse,
+    Usage,
+)
+from .scoring import (
+    choice_confidence,
+    combine_escape,
+    expected_score,
+    normalise_levels,
+    restricted_softmax,
+    score_confidence,
+)
+from .vision import extract_image
+
+# The label prior is a property of the token and how many labels compete, so
+# branches sharing a shape share a bucket. A score's letter names a fixed rung, so a
+# score never shares a bucket with a choice of the same width.
+_PRIOR_KIND = {
+    "choice": "choice",
+    "member": "choice",
+    "score_level": "binary",
+    "score_joint": "score",
+    "noul": "binary",
+}
+
+
+@dataclass
+class Diagnostics:
+    """Per-branch signals that Jev does not return, surfaced on request."""
+
+    candidate_mass: dict[str, list[float]] = field(default_factory=dict)
+    prior_applied: dict[str, bool] = field(default_factory=dict)
+    branch_counts: dict[str, int] = field(default_factory=dict)
+
+
+def load_model(model_id: str | None = None, config: Config | None = None) -> Backend:
+    """Load a model through transformers and return it as a Backend.
+
+    `model_id` names a Hugging Face repo or a local directory, and overrides the one on
+    `config`. Where the weights land is `config.models_dir`. A ComfyUI node skips this
+    entirely and passes its own Backend over a model the workflow already loaded.
+
+    The import sits inside the call so the core install needs no torch.
+    """
+    config = config if config is not None else Config()
+    if model_id is not None:
+        config = replace(config, model_id=model_id)
+
+    try:
+        from .backends.hf import HFBackend
+    except ImportError as error:
+        raise MissingDependencyError(
+            "no backend was passed and the local transformers backend is not installed. "
+            'Install it with: pip install "logit-classifier[hf]"'
+        ) from error
+    return HFBackend(config)
+
+
+def _model_identity(backend: Backend) -> str:
+    """Return the id the fitted temperature and the prior fingerprint key on."""
+    return getattr(backend, "canonical_model_id", None) or backend.model_id
+
+
+class Classifier:
+    """Turns one System One request into one calibrated answer per question.
+
+    A caller that already holds a loaded model, such as a ComfyUI node over a
+    resident text encoder, passes its own backend and no weights are loaded here.
+
+    One instance belongs to one backend. The fitted temperature and the prior bucket
+    are both derived from the backend's model at construction, so swapping the
+    backend afterwards leaves both belonging to the previous model.
+    """
+
+    def __init__(self, config: Config, backend: Backend | None = None) -> None:
+        self.config = config
+        self.backend = backend if backend is not None else load_model(config=config)
+        # Keyed on the model the backend actually loaded, not on a Config field a
+        # caller passing its own backend never had reason to set.
+        self.temperature = (
+            config.temperature if config.temperature is not None
+            else fitted_temperature(_model_identity(self.backend))
+        )
+        self.priors = PriorStore(config.calibration_path, self._fingerprint())
+
+    def _fingerprint(self) -> str:
+        """Identify the prompt shape a stored prior was measured under."""
+        parts = [_model_identity(self.backend), ANSWER_PREFILL, SYSTEM_PROMPT, self.config.score_method,
+                 str(PROMPT_VERSION), str(self.config.abstain)]
+        return hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()[:16]
+
+    def persist_prior(self) -> None:
+        """Write the running prior, unless the caller turned the prior off.
+
+        The service calls this again at shutdown, so the gate lives here rather than
+        at each call site.
+        """
+        if self.config.use_prior_debias:
+            self.priors.save()
+
+    def _calibrated(self, branch: Branch, logits: BranchLogits) -> tuple[np.ndarray, bool]:
+        """Raw distribution feeds the prior estimate, calibrated one is returned."""
+        kind = _PRIOR_KIND[branch.kind]
+        # The escape label means "none of these" wherever it sits, so its mass is content
+        # rather than letter bias, and only the lettered options get a prior.
+        lettered = branch.label_count - 1 if branch.kind == "member" else branch.label_count
+        letters_prior: np.ndarray | None = None
+        log_prior: np.ndarray | None = None
+
+        # One switch turns learning, application and persistence off together, so a host
+        # that asked for no debias gets the same numbers on every call.
+        if self.config.use_prior_debias:
+            self.priors.observe(kind, restricted_softmax(logits.z[:lettered]))
+            letters_prior = self.priors.log_prior(kind, lettered)
+        if letters_prior is not None:
+            log_prior = np.zeros(branch.label_count, dtype=np.float64)
+            log_prior[:lettered] = letters_prior
+        calibrated = restricted_softmax(logits.z, log_prior, self.temperature)
+        return calibrated, log_prior is not None
+
+    def _letterings(self, questions: dict[str, Question]) -> list[dict[str, Question]]:
+        """One question map per lettering, the first in the order the caller gave.
+
+        Only choice options are reordered. A score lists its levels lowest first, so
+        that order carries meaning, and a noul has two fixed sides. Each lettering
+        draws from a fixed seed, so the same request still answers the same way.
+        """
+        rounds = [questions]
+
+        for seed in range(1, max(1, self.config.permutations)):
+            shuffler = random.Random(seed)
+            reordered: dict[str, Question] = {}
+            for qid, question in questions.items():
+                if not isinstance(question, ChoiceQuestion):
+                    continue
+                names = list(question.criteria)
+                shuffler.shuffle(names)
+                criteria = {name: question.criteria[name] for name in names}
+                reordered[qid] = replace(question, criteria=criteria)
+            if reordered:
+                rounds.append(reordered)
+        return rounds
+
+    def _suffix_ids(self, state: object, branches: list[Branch], has_image: bool,
+                    prefix_text: str) -> list[list[int]]:
+        """Each branch's tokens after the shared prefix, which it must begin with."""
+        ids: list[list[int]] = []
+
+        for branch in branches:
+            rendered = self.backend.render(
+                SYSTEM_PROMPT, branch_content(state, branch, has_image), ANSWER_PREFILL
+            )
+            if not rendered.startswith(prefix_text):
+                raise BackendContractError(
+                    f"{type(self.backend).__name__}.render did not begin the closed render "
+                    f"of question {branch.question_id!r} with the open-ended render of the "
+                    "same state, so the shared prefix cannot be split off by offset"
+                )
+            ids.append(self.backend.encode(rendered[len(prefix_text):]))
+        return ids
+
+    def classify(
+        self, request: SystemOneRequest, *, allow_image_paths: bool = True
+    ) -> tuple[SystemOneResponse, Diagnostics]:
+        rounds = self._letterings(request.questions)
+        diagnostics = Diagnostics()
+        probabilities: list[np.ndarray] = []
+        backend_name = type(self.backend).__name__
+        drafts: list[dict[str, Answer]] = []
+        start = 0
+
+        state, image = extract_image(request.state, allow_paths=allow_image_paths)
+        seen = image is not None
+        prefix_text = self.backend.render(
+            SYSTEM_PROMPT, prefix_content(state, seen), ANSWER_PREFILL, open_ended=True
+        )
+        prefix_ids, vision = self.backend.encode_prefix(prefix_text, image)
+        round_branches = [build_branches(q, self.config.score_method, self.config.abstain)
+                          for q in rounds]
+        branches: list[Branch] = [b for lettering in round_branches for b in lettering]
+        suffix_ids = self._suffix_ids(state, branches, seen, prefix_text)
+        scored = self.backend.score(prefix_ids, suffix_ids,
+                                    [b.label_count for b in branches], vision)
+
+        if len(scored) != len(branches):
+            raise BackendContractError(
+                f"{backend_name}.score returned {len(scored)} rows for "
+                f"{len(branches)} branches"
+            )
+        for branch, logits in zip(branches, scored, strict=True):
+            if len(logits.z) != branch.label_count:
+                raise BackendContractError(
+                    f"{backend_name}.score returned a {len(logits.z)}-wide row for "
+                    f"question {branch.question_id}, which needs {branch.label_count}"
+                )
+            calibrated, applied = self._calibrated(branch, logits)
+            probabilities.append(calibrated)
+            diagnostics.candidate_mass.setdefault(branch.question_id, []).append(
+                round(logits.candidate_mass, 6)
+            )
+            diagnostics.prior_applied[branch.question_id] = applied
+            diagnostics.branch_counts[branch.question_id] = (
+                diagnostics.branch_counts.get(branch.question_id, 0) + 1
+            )
+        self.persist_prior()
+
+        for questions, lettering in zip(rounds, round_branches, strict=True):
+            window = slice(start, start + len(lettering))
+            start += len(lettering)
+            drafts.append(self._round_answers(questions, lettering, probabilities[window]))
+        answers = self._merge(request.questions, drafts)
+
+        usage = Usage(
+            input_tokens=len(prefix_ids) + sum(len(s) for s in suffix_ids),
+            output_tokens=len(branches),
+        )
+        response = SystemOneResponse(
+            model=self.config.served_model_id, answers=answers, usage=usage
+        )
+        return response, diagnostics
+
+    def _round_answers(self, questions: dict[str, Question], branches: list[Branch],
+                       probabilities: list[np.ndarray]) -> dict[str, Answer]:
+        answers: dict[str, Answer] = {}
+
+        for qid, question in questions.items():
+            span = [i for i, b in enumerate(branches) if b.question_id == qid]
+            answers[qid] = self._assemble(question, [branches[i] for i in span],
+                                          [probabilities[i] for i in span])
+        return answers
+
+    def _merge(self, questions: dict[str, Question],
+               drafts: list[dict[str, Answer]]) -> dict[str, Answer]:
+        """Average each option's probability across the letterings that scored it."""
+        answers: dict[str, Answer] = {}
+
+        for qid in questions:
+            parts = [draft[qid] for draft in drafts if qid in draft]
+            if len(parts) == 1:
+                answers[qid] = parts[0]
+                continue
+            answers[qid] = self._average_choice([p for p in parts if isinstance(p, ChoiceAnswer)])
+        return answers
+
+    def _average_choice(self, parts: list[ChoiceAnswer]) -> ChoiceAnswer:
+        names = list(parts[0].probabilities)
+        stacked = np.array([[part.probabilities[name] for name in names] for part in parts])
+        mean = stacked.mean(axis=0)
+        mean = mean / mean.sum()
+        declined = [part.abstain for part in parts if part.abstain is not None]
+        return ChoiceAnswer(
+            choice=names[int(np.argmax(mean))],
+            confidence=round(choice_confidence(mean), 6),
+            probabilities={name: round(float(p), 6) for name, p in zip(names, mean, strict=True)},
+            abstain=round(float(np.mean(declined)), 6) if declined else None,
+        )
+
+    def _assemble(self, question: Question, branches: list[Branch],
+                  probabilities: list[np.ndarray]) -> Answer:
+        if isinstance(question, ChoiceQuestion):
+            return self._choice_answer(question, branches, probabilities)
+        if isinstance(question, ScoreQuestion):
+            return self._score_answer(question, branches, probabilities)
+        if isinstance(question, NoulQuestion):
+            return NoulAnswer(noul=round(float(probabilities[0][0]), 6))
+        raise TypeError(f"unsupported question type {type(question)!r}")
+
+    def _choice_answer(self, question: ChoiceQuestion, branches: list[Branch],
+                       probabilities: list[np.ndarray]) -> ChoiceAnswer:
+        names = list(question.criteria)
+        declined: float | None = None
+
+        if branches[0].kind == "choice":
+            flat = probabilities[0]
+        else:
+            flat, declined = combine_escape(probabilities)
+        mapping = {name: round(float(p), 6) for name, p in zip(names, flat, strict=True)}
+        winner = names[int(np.argmax(flat))]
+        return ChoiceAnswer(
+            choice=winner,
+            confidence=round(choice_confidence(flat), 6),
+            probabilities=mapping,
+            abstain=None if declined is None else round(declined, 6),
+        )
+
+    def _score_answer(self, question: ScoreQuestion, branches: list[Branch],
+                      probabilities: list[np.ndarray]) -> ScoreAnswer:
+        if branches[0].kind == "score_joint":
+            levels = probabilities[0]
+        else:
+            # Each level was judged alone, so index 0 of each branch is its yes.
+            levels = normalise_levels(np.array([float(p[0]) for p in probabilities]))
+        return ScoreAnswer(
+            score=round(expected_score(levels), 6),
+            confidence=round(score_confidence(levels), 6),
+            legend={str(i): level for i, level in enumerate(question.criteria)},
+            probabilities={str(i): round(float(p), 6) for i, p in enumerate(levels)},
+        )

logit_classifier/cli.py

@@ -0,0 +1,51 @@
+"""The `logit-classifier` command.
+
+`serve` is the front door for anyone who installed from PyPI rather than cloning,
+since it needs no uvicorn invocation and no import path.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from . import __version__
+from .config import Config
+from .deps import require
+
+
+def _serve(host: str, port: int, reload: bool) -> int:
+    # uvicorn imports the app from a string, so a missing fastapi surfaces as a raw
+    # traceback out of its importer rather than as the line that installs the extra.
+    require("fastapi", "service")
+    uvicorn = require("uvicorn", "service")
+    uvicorn.run("logit_classifier.service:app", host=host, port=port, reload=reload)
+    return 0
+
+
+def _show_config() -> int:
+    config = Config.from_env()
+    for name, value in vars(config).items():
+        print(f"{name}={value}")
+    return 0
+
+
+def _parser(prog: str) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog=prog, description=__doc__.splitlines()[0])
+    parser.add_argument("--version", action="version", version=__version__)
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    serve = commands.add_parser("serve", help="run the HTTP service and the browser page")
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", type=int, default=8077)
+    serve.add_argument("--reload", action="store_true", help="restart on a source change")
+
+    commands.add_parser("config", help="print the configuration the LOGIT_ variables produce")
+    return parser
+
+
+def main(argv: list[str] | None = None, prog: str = "logit-classifier") -> int:
+    args = _parser(prog).parse_args(argv)
+
+    if args.command == "serve":
+        return _serve(args.host, args.port, args.reload)
+    return _show_config()