logit-classifier 0.1.0__py3-none-any.whl

logit_classifier/schema.py

@@ -0,0 +1,285 @@
+"""Request and response models mirroring the TypeSafe System One wire format.
+
+Field names, nesting and answer shapes follow the published Jev contract so a
+Jev request body is accepted unchanged. Deviations are noted where they occur.
+
+Plain dataclasses rather than pydantic, so a ComfyUI node pack importing this
+package never pulls a pydantic version into ComfyUI's own venv.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from typing import Any, Literal
+
+from .config import JEV_MAX_OPTIONS, JEV_MAX_QUESTIONS, JEV_MAX_SCORE_LEVELS
+from .errors import LogitClassifierError
+
+# Jev accepts a string, object or array anywhere prose is expected.
+JSONContent = str | dict[str, Any] | list[Any]
+
+DEFAULT_MODEL = "logit-latest"
+
+_REQUEST_FIELDS = frozenset({"state", "model", "questions"})
+_QUESTION_FIELDS = frozenset({"type", "instructions", "criteria"})
+_NOUL_CRITERIA_FIELDS = frozenset({"true", "false"})
+_QUESTION_TYPES = ("choice", "score", "noul")
+
+
+class SchemaError(LogitClassifierError, ValueError):
+    """A body does not match the System One contract.
+
+    `field` carries the dotted path to the offending value, which is what the
+    service reports back so a caller can find it without guessing.
+    """
+
+    def __init__(self, message: str, field: str | None = None) -> None:
+        super().__init__(message)
+        self.message = message
+        self.field = field or None
+
+
+def render_content(value: JSONContent | None) -> str:
+    """Flatten a JSONContent value into the text the model sees."""
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, indent=2, ensure_ascii=False, sort_keys=True)
+
+
+@dataclass(frozen=True, kw_only=True)
+class NoulCriteria:
+    """Wording for the two sides of a noul. `true` and `false` are reserved words."""
+
+    true_: JSONContent | None = None
+    false_: JSONContent | None = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class ChoiceQuestion:
+    """Pick one of the named options."""
+
+    type: Literal["choice"] = "choice"
+    instructions: JSONContent | None = None
+    criteria: dict[str, JSONContent | None]
+
+
+@dataclass(frozen=True, kw_only=True)
+class ScoreQuestion:
+    """Place the state on an ordered scale, lowest level first."""
+
+    type: Literal["score"] = "score"
+    instructions: JSONContent | None = None
+    criteria: list[JSONContent]
+
+
+@dataclass(frozen=True, kw_only=True)
+class NoulQuestion:
+    """How true one statement is of the state."""
+
+    type: Literal["noul"] = "noul"
+    instructions: JSONContent | None = None
+    criteria: NoulCriteria | None = None
+
+
+Question = ChoiceQuestion | ScoreQuestion | NoulQuestion
+
+
+@dataclass(frozen=True, kw_only=True)
+class SystemOneRequest:
+    """One state, and every question asked about it."""
+
+    state: JSONContent
+    model: str = DEFAULT_MODEL
+    questions: dict[str, Question]
+
+
+# Answer field order matches the published response examples, and asdict follows
+# the order the fields are declared in.
+@dataclass(frozen=True, kw_only=True)
+class ChoiceAnswer:
+    type: Literal["choice"] = "choice"
+    choice: str
+    confidence: float
+    probabilities: dict[str, float]
+    # How strongly the model wants none of the offered options. It sits outside
+    # probabilities, which still covers the caller's own options and sums to 1.
+    abstain: float | None = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class ScoreAnswer:
+    type: Literal["score"] = "score"
+    score: float
+    confidence: float
+    legend: dict[str, JSONContent]
+    probabilities: dict[str, float]
+
+
+@dataclass(frozen=True, kw_only=True)
+class NoulAnswer:
+    type: Literal["noul"] = "noul"
+    noul: float
+
+
+Answer = ChoiceAnswer | ScoreAnswer | NoulAnswer
+
+
+@dataclass(frozen=True, kw_only=True)
+class Usage:
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class SystemOneResponse:
+    model: str
+    answers: dict[str, Answer]
+    usage: Usage
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass(frozen=True, kw_only=True)
+class ErrorBody:
+    """The error shape this service returns.
+
+    Jev documents its status codes but publishes no error body, so this shape is ours.
+    """
+
+    error: str
+    message: str
+    field: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def _path(*parts: str) -> str:
+    return ".".join(part for part in parts if part)
+
+
+def _require_mapping(value: Any, where: str) -> dict[str, Any]:
+    if not isinstance(value, dict):
+        raise SchemaError(f"expected an object, got {type(value).__name__}", where)
+    for key in value:
+        if not isinstance(key, str):
+            raise SchemaError(f"object keys must be strings, got {type(key).__name__}", where)
+    return value
+
+
+def _reject_unknown(mapping: dict[str, Any], allowed: frozenset[str], where: str) -> None:
+    unknown = sorted(set(mapping) - allowed)
+    if unknown:
+        raise SchemaError(f"unexpected field {unknown[0]!r}", _path(where, unknown[0]))
+
+
+def _content(value: Any, where: str) -> JSONContent:
+    if isinstance(value, str | dict | list):
+        return value
+    raise SchemaError(f"expected a string, object or array, got {type(value).__name__}", where)
+
+
+def _optional_content(value: Any, where: str) -> JSONContent | None:
+    return None if value is None else _content(value, where)
+
+
+def _choice_criteria(value: Any, where: str) -> dict[str, JSONContent | None]:
+    mapping = _require_mapping(value, where)
+
+    if not 2 <= len(mapping) <= JEV_MAX_OPTIONS:
+        raise SchemaError(
+            f"a choice needs 2 to {JEV_MAX_OPTIONS} options, got {len(mapping)}", where
+        )
+    return {name: _optional_content(body, _path(where, name)) for name, body in mapping.items()}
+
+
+def _score_criteria(value: Any, where: str) -> list[JSONContent]:
+    if not isinstance(value, list):
+        raise SchemaError(f"expected an array, got {type(value).__name__}", where)
+    if not 2 <= len(value) <= JEV_MAX_SCORE_LEVELS:
+        raise SchemaError(
+            f"a score needs 2 to {JEV_MAX_SCORE_LEVELS} levels, got {len(value)}", where
+        )
+    return [_content(level, _path(where, str(index))) for index, level in enumerate(value)]
+
+
+def _noul_criteria(value: Any, where: str) -> NoulCriteria:
+    mapping = _require_mapping(value, where)
+
+    _reject_unknown(mapping, _NOUL_CRITERIA_FIELDS, where)
+    return NoulCriteria(
+        true_=_optional_content(mapping.get("true"), _path(where, "true")),
+        false_=_optional_content(mapping.get("false"), _path(where, "false")),
+    )
+
+
+def _parse_question(value: Any, where: str) -> Question:
+    mapping = _require_mapping(value, where)
+    kind = mapping.get("type")
+    instructions = _optional_content(mapping.get("instructions"), _path(where, "instructions"))
+    criteria = mapping.get("criteria")
+    criteria_at = _path(where, "criteria")
+
+    if kind not in _QUESTION_TYPES:
+        raise SchemaError(
+            f"type must be one of {', '.join(_QUESTION_TYPES)}, got {kind!r}",
+            _path(where, "type"),
+        )
+    _reject_unknown(mapping, _QUESTION_FIELDS, where)
+
+    if kind == "choice":
+        if criteria is None:
+            raise SchemaError("a choice needs criteria", criteria_at)
+        return ChoiceQuestion(
+            instructions=instructions, criteria=_choice_criteria(criteria, criteria_at)
+        )
+    if kind == "score":
+        if criteria is None:
+            raise SchemaError("a score needs criteria", criteria_at)
+        return ScoreQuestion(
+            instructions=instructions, criteria=_score_criteria(criteria, criteria_at)
+        )
+    return NoulQuestion(
+        instructions=instructions,
+        criteria=None if criteria is None else _noul_criteria(criteria, criteria_at),
+    )
+
+
+def parse_questions(value: Any) -> dict[str, Question]:
+    """Validate a question map on its own, for a caller holding no full request."""
+    mapping = _require_mapping(value, "questions")
+
+    if not 1 <= len(mapping) <= JEV_MAX_QUESTIONS:
+        raise SchemaError(
+            f"1 to {JEV_MAX_QUESTIONS} questions are allowed, got {len(mapping)}", "questions"
+        )
+    return {qid: _parse_question(body, _path("questions", qid)) for qid, body in mapping.items()}
+
+
+def parse_request(body: Any) -> SystemOneRequest:
+    """Validate a decoded JSON body and return the request it describes.
+
+    Every failure raises SchemaError carrying the dotted path to the value that
+    failed, so the first rejection names its own field rather than the body.
+    """
+    mapping = _require_mapping(body, "")
+
+    _reject_unknown(mapping, _REQUEST_FIELDS, "")
+    if "state" not in mapping:
+        raise SchemaError("state is required", "state")
+    if "questions" not in mapping:
+        raise SchemaError("questions is required", "questions")
+
+    model = mapping.get("model", DEFAULT_MODEL)
+    if not isinstance(model, str):
+        raise SchemaError(f"expected a string, got {type(model).__name__}", "model")
+
+    return SystemOneRequest(
+        state=_content(mapping["state"], "state"),
+        model=model,
+        questions=parse_questions(mapping["questions"]),
+    )

logit_classifier/scoring.py

@@ -0,0 +1,102 @@
+"""Turning label logits into calibrated probabilities."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def restricted_softmax(
+    z: np.ndarray, log_prior: np.ndarray | None = None, temperature: float = 1.0
+) -> np.ndarray:
+    """Normalise over the option labels only, in float64.
+
+    Subtracting a log-prior is the shared operation behind batch calibration,
+    contextual calibration and PriDe. Temperature then widens the distribution,
+    which a restricted softmax over an aligned model is otherwise too sharp for.
+    """
+    adjusted = z.astype(np.float64)
+
+    if log_prior is not None:
+        adjusted = adjusted - log_prior
+    if temperature and temperature != 1.0:
+        adjusted = adjusted / temperature
+    adjusted = adjusted - adjusted.max()
+    weights: np.ndarray = np.exp(adjusted)
+    normalised: np.ndarray = weights / weights.sum()
+    return normalised
+
+
+def choice_confidence(probabilities: np.ndarray) -> float:
+    """Peak probability rescaled so uniform maps to 0 and a point mass to 1.
+
+    Ported with score_confidence from TypeSafe's own
+    system_one_adapter._utils.confidence_metrics, so a System One client reading
+    our answers computes the same number it would from theirs.
+    """
+    count = len(probabilities)
+
+    if count < 2:
+        return 1.0
+    peak = float(probabilities.max())
+    return float(np.clip((count * peak - 1.0) / (count - 1.0), 0.0, 1.0))
+
+
+def score_confidence(probabilities: np.ndarray) -> float:
+    """Concentration of an ordered distribution around its modal level.
+
+    Levels are ordinal, so mass one level from the mode counts less than mass
+    four levels away. The choice formula cannot express that, and it reports a
+    tight two level split as low confidence.
+    """
+    count = len(probabilities)
+
+    if count < 2:
+        return 1.0
+    indices = np.arange(count, dtype=np.float64)
+    spread = float(probabilities @ np.abs(indices - int(np.argmax(probabilities))))
+    uniform_spread = float(np.abs(indices - (count - 1) / 2.0).mean())
+    return float(max(0.0, 1.0 - spread / uniform_spread))
+
+
+def normalise_levels(yes_probabilities: np.ndarray) -> np.ndarray:
+    """Turn independent per-level yes probabilities into one distribution."""
+    total = float(yes_probabilities.sum())
+
+    if total <= 0.0:
+        return np.full(len(yes_probabilities), 1.0 / len(yes_probabilities))
+    return yes_probabilities / total
+
+
+def expected_score(level_probabilities: np.ndarray) -> float:
+    """Probability-weighted mean of the zero-based level indices."""
+    return float(level_probabilities @ np.arange(len(level_probabilities), dtype=np.float64))
+
+
+def combine_escape(group_probabilities: list[np.ndarray]) -> tuple[np.ndarray, float]:
+    """Merge groups whose last label is "none of these" into one distribution.
+
+    A group's weight is the mass it did not put on its escape label, so the
+    groups compete without ever being summarised for the model. That removes the
+    separate group branch, and with it the label bias that branch carried.
+
+    The second return is the abstain score, which is how hard the least declining
+    group declined. Every group has to decline for the question to be unanswerable,
+    so the weakest decline is the binding one. It is reported beside the distribution
+    rather than inside it.
+    """
+    weights = np.array([1.0 - float(p[-1]) for p in group_probabilities], dtype=np.float64)
+    parts: list[np.ndarray] = []
+
+    if weights.sum() <= 0:
+        weights = np.ones(len(group_probabilities), dtype=np.float64)
+    weights = weights / weights.sum()
+
+    for weight, probabilities in zip(weights, group_probabilities, strict=True):
+        members = probabilities[:-1]
+        total = members.sum()
+        share = members / total if total > 0 else np.full(len(members), 1.0 / len(members))
+        parts.append(weight * share)
+    combined = np.concatenate(parts)
+    total = combined.sum()
+    declined = float(min(float(p[-1]) for p in group_probabilities))
+    return (combined / total if total > 0 else combined), declined

logit_classifier/service.py

@@ -0,0 +1,174 @@
+"""HTTP service exposing the Jev-shaped System One endpoint."""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Any
+
+from fastapi import Body, FastAPI, Query, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
+
+from . import __version__
+from .backends.base import VisionUnsupportedError
+from .classifier import Classifier
+from .config import Config
+from .errors import LogitClassifierError
+from .schema import ErrorBody, SchemaError, parse_request
+from .vision import ImageError
+
+WEB_DIR = Path(__file__).parent / "web"
+
+_log = logging.getLogger(__name__)
+
+EXAMPLE_REQUEST = {
+    "state": "The integration keeps failing and I am losing sales.",
+    "questions": {
+        "department": {
+            "type": "choice",
+            "instructions": "Which team should handle this",
+            "criteria": {"billing": "Payment issues", "technical": "Integration problems"},
+        }
+    },
+}
+
+
+@dataclass(frozen=True)
+class Loaded:
+    """What the lifespan built, present only once the model is resident."""
+
+    config: Config
+    classifier: Classifier
+    load_seconds: float
+
+
+_loaded: Loaded | None = None
+# One GPU, and batch composition must stay a pure function of the request, so
+# requests are served one at a time.
+_gpu_lock = threading.Lock()
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+    global _loaded
+    config = Config.from_env()
+    started = time.time()
+    classifier = Classifier(config)
+    _loaded = Loaded(config=config, classifier=classifier,
+                     load_seconds=round(time.time() - started, 2))
+    yield
+    classifier.persist_prior()
+
+
+app = FastAPI(title="Logit Classifier", version=__version__, lifespan=lifespan)
+
+
+def _rejected(message: str, field: str | None = None) -> JSONResponse:
+    body = ErrorBody(error="unprocessable_entity", message=message, field=field)
+    return JSONResponse(status_code=422, content=body.to_dict())
+
+
+def _failed(message: str) -> JSONResponse:
+    body = ErrorBody(error="internal_error", message=message)
+    return JSONResponse(status_code=500, content=body.to_dict())
+
+
+def _unavailable() -> JSONResponse:
+    body = ErrorBody(error="service_unavailable", message="model is still loading")
+    return JSONResponse(status_code=503, content=body.to_dict())
+
+
+@app.exception_handler(RequestValidationError)
+async def on_validation_error(_request: Request, exc: RequestValidationError) -> JSONResponse:
+    first = exc.errors()[0] if exc.errors() else {}
+    location = ".".join(str(p) for p in first.get("loc", []) if p != "body")
+    return _rejected(first.get("msg", "request failed validation"), location or None)
+
+
+@app.post("/v1/systemone")
+def system_one(
+    body: Annotated[dict[str, Any], Body(openapi_examples={"jev": {"value": EXAMPLE_REQUEST}})],
+    diagnostics: Annotated[bool, Query()] = False,
+) -> Any:
+    loaded = _loaded
+
+    if loaded is None:
+        return _unavailable()
+    try:
+        request = parse_request(body)
+    except SchemaError as error:
+        return _rejected(error.message, error.field)
+
+    started = time.time()
+    try:
+        with _gpu_lock:
+            response, diag = loaded.classifier.classify(request, allow_image_paths=False)
+    except (ImageError, VisionUnsupportedError) as error:
+        return _rejected(str(error), "state")
+    except LogitClassifierError as error:
+        # A backend contract or label failure is ours, not the caller's, so it must
+        # not come back as a rejected request they could fix. The traceback is the
+        # operator's only diagnostic, and returning a body swallows it.
+        _log.exception("classify failed")
+        return _failed(str(error))
+    except Exception as error:
+        # A CUDA out of memory or a kernel error must still reach a Jev caller as an
+        # ErrorBody, never as a plain-text 500.
+        _log.exception("classify failed")
+        return _failed(f"{type(error).__name__}: {error}")
+    payload = response.to_dict()
+
+    if diagnostics:
+        payload["diagnostics"] = {
+            "candidate_mass": diag.candidate_mass,
+            "prior_applied": diag.prior_applied,
+            "branch_counts": diag.branch_counts,
+            "latency_ms": round((time.time() - started) * 1000, 1),
+            "temperature": loaded.classifier.temperature,
+            "score_method": loaded.config.score_method,
+            "prior_observations": loaded.classifier.priors.stats(),
+        }
+    return payload
+
+
+@app.get("/v1/models")
+def models() -> Any:
+    loaded = _loaded
+
+    if loaded is None:
+        return _unavailable()
+    return {
+        "models": [
+            {
+                "name": loaded.config.served_model_id,
+                "description": f"Restricted-label scorer over {loaded.config.model_id}",
+                "aliases": list(loaded.config.model_aliases),
+            }
+        ]
+    }
+
+
+@app.get("/healthz")
+def healthz() -> dict[str, Any]:
+    loaded = _loaded
+    return {
+        "ready": loaded is not None,
+        "model_id": loaded.config.model_id if loaded else None,
+        "sees_images": loaded.classifier.backend.sees_images if loaded else None,
+        "load_seconds": loaded.load_seconds if loaded else None,
+    }
+
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(WEB_DIR / "index.html")
+
+
+app.mount("/static", StaticFiles(directory=WEB_DIR), name="static")

logit_classifier/vision.py

@@ -0,0 +1,95 @@
+"""Pulling an image out of the request state.
+
+The state stays JSON, so an image arrives as a data URL, a bare base64 string or a
+path. It is lifted out before the state is rendered as text, because base64 in the
+prompt would cost thousands of tokens and say nothing.
+"""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import io
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+from .deps import require
+from .errors import LogitClassifierError
+
+if TYPE_CHECKING:
+    from PIL.Image import Image
+
+IMAGE_KEYS = ("image", "screenshot")
+
+
+class ImageError(LogitClassifierError, ValueError):
+    """The state named an image that could not be read."""
+
+
+def _from_base64(value: str) -> Image:
+    pil_image = require("PIL.Image", "service")
+
+    try:
+        raw = base64.b64decode(value, validate=True)
+    except (binascii.Error, ValueError) as error:
+        raise ImageError(f"image is neither a readable path nor valid base64: {error}") from error
+    try:
+        return cast("Image", pil_image.open(io.BytesIO(raw)).convert("RGB"))
+    except (OSError, ValueError, pil_image.DecompressionBombError) as error:
+        raise ImageError(f"decoded bytes are not a readable image: {error}") from error
+
+
+def _decode(value: str, key: str, *, allow_paths: bool) -> Image:
+    """Decode a data URL, a path to a file, or bare base64, by trying rather than guessing.
+
+    Length cannot tell a path from base64, since a small PNG encodes shorter than many
+    paths. Whether the file exists can.
+    """
+    path: Path | None = None
+    pil_image = require("PIL.Image", "service")
+
+    if value.startswith("data:"):
+        _, _, payload = value.partition(",")
+        if not payload:
+            raise ImageError("data URL carried no payload")
+        return _from_base64(payload)
+    # A path probe over HTTP would let a client read server files, and a UNC path would
+    # send the service account's credentials to a host the client names.
+    if not allow_paths:
+        try:
+            return _from_base64(value)
+        except ImageError as error:
+            raise ImageError(
+                f"state.{key} is not a data URL or base64 image. The HTTP service accepts "
+                f"only those two forms: {error}"
+            ) from error
+    try:
+        candidate = Path(value)
+        path = candidate if candidate.is_file() else None
+    except (OSError, ValueError):
+        path = None  # too long or illegal as a path, so it is not one
+    if path is None:
+        return _from_base64(value)
+    # UnidentifiedImageError subclasses OSError, so reading the file stays outside
+    # the probe above, which would report a real file as bad base64.
+    try:
+        opened = pil_image.open(path).convert("RGB")
+    except (OSError, ValueError, pil_image.DecompressionBombError) as error:
+        raise ImageError(f"{path} is not a readable image: {error}") from error
+    return cast("Image", opened)
+
+
+def extract_image(state: Any, *, allow_paths: bool = True) -> tuple[Any, Image | None]:
+    """Return the state with any image removed, and the image itself.
+
+    A state that is not a mapping carries no image, so it comes back untouched.
+    """
+    if not isinstance(state, dict):
+        return state, None
+
+    for key in IMAGE_KEYS:
+        value = state.get(key)
+        if isinstance(value, str) and value:
+            remainder = {k: v for k, v in state.items() if k != key}
+            return (remainder or "(see image)"), _decode(value, key, allow_paths=allow_paths)
+    return state, None