logit-classifier 0.1.0__py3-none-any.whl

logit_classifier/config.py

@@ -0,0 +1,187 @@
+"""Runtime configuration for the classifier service."""
+
+from __future__ import annotations
+
+import math
+import os
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+from .errors import ConfigError
+
+# Jev caps a Choice at 255 options and a Score at 10 levels. Matching those caps
+# keeps a Jev-shaped request valid here without translation.
+JEV_MAX_OPTIONS = 255
+JEV_MAX_SCORE_LEVELS = 10
+JEV_MAX_QUESTIONS = 256
+
+# The tokenizer merges a letter into a preceding space, so the prefill must end
+# on a non-space. Verified against Qwen: "Answer: (" ends on the token "Ġ(" and
+# "Ġ(A" is absent from the vocabulary, leaving each letter its own token.
+ANSWER_PREFILL = "Answer: ("
+
+# Temperature belongs to the model, not to the method. Each of these was fitted on the
+# held out half of a 924 row Banking77 sample and chosen on calibration error. Accuracy
+# also climbs with temperature above 52 options, but that is the split path artefact
+# rather than a gain, so the fit does not follow it.
+FITTED_TEMPERATURES = {
+    "Qwen/Qwen3-VL-4B-Instruct": 1.25,
+    "Qwen/Qwen3-4B-Instruct-2507": 6.0,
+}
+
+# What an unfitted model gets. Published fitted temperatures for this readout cluster
+# near 1 in domain and near 3 out of domain, so the midpoint is the least wrong guess.
+DEFAULT_TEMPERATURE = 2.5
+
+
+def fitted_temperature(model_id: str) -> float:
+    return FITTED_TEMPERATURES.get(model_id, DEFAULT_TEMPERATURE)
+
+
+def canonical_model_id(model_id: str) -> str:
+    """Map a local directory holding shipped weights onto its Hugging Face repo id.
+
+    A local path would otherwise miss FITTED_TEMPERATURES and fall back to the default.
+    """
+    if model_id in FITTED_TEMPERATURES:
+        return model_id
+    folder = model_id.replace("\\", "/").rstrip("/").rsplit("/", 1)[-1]
+    for known in FITTED_TEMPERATURES:
+        if known.rsplit("/", 1)[-1] == folder:
+            return known
+    return model_id
+
+
+# Offering an escape label costs nothing and scores how badly the model wants to decline.
+ABSTAIN = True
+
+# The two branch layouts _score_branches implements. Any other spelling would reach
+# the independent path without a word.
+SCORE_METHODS = frozenset({"joint", "independent"})
+
+# Below this many observations the running prior is noise, so it stays unused.
+PRIOR_MIN_OBSERVATIONS = 32
+
+# A label the model never picks drives its mean toward zero, and subtracting the log of a
+# near-zero prior would hand that label an unbounded boost. Blending toward uniform bounds
+# the correction. Measured on Banking77, an unsmoothed prior reached a 40 million to one
+# spread across 40 labels.
+PRIOR_SMOOTHING = 0.1
+
+
+def _reject_score_method(where: str, value: str) -> None:
+    if value not in SCORE_METHODS:
+        raise ConfigError(
+            f"{where}={value!r} is not a score method. "
+            f"Accepted values are {', '.join(sorted(SCORE_METHODS))}."
+        )
+
+
+def _reject_temperature(where: str, value: float | None) -> None:
+    if value is not None and not (math.isfinite(value) and value > 0):
+        raise ConfigError(
+            f"{where}={value!r} is not a temperature. "
+            "Accepted values are None or a finite number above 0."
+        )
+
+
+def _reject_below_one(where: str, value: int) -> None:
+    if value < 1:
+        raise ConfigError(
+            f"{where}={value!r} is below 1. Accepted values are whole numbers of 1 or more."
+        )
+
+
+def _env_switch(name: str, default: bool) -> bool:
+    raw = os.environ.get(name)
+
+    if raw is None:
+        return default
+    if raw not in ("0", "1"):
+        raise ConfigError(f"{name}={raw!r} is not a switch. Accepted values are 0 and 1.")
+    return raw == "1"
+
+
+def _env_number[N: (int, float)](name: str, parse: Callable[[str], N]) -> N | None:
+    raw = os.environ.get(name)
+
+    if not raw:
+        return None
+    try:
+        return parse(raw)
+    except ValueError as error:
+        raise ConfigError(
+            f"{name}={raw!r} is not a number. Accepted values are {parse.__name__} literals."
+        ) from error
+
+
+@dataclass(frozen=True)
+class Config:
+    # Reads images as well as text, and scored 0.626 against 0.554 for the text only
+    # Qwen3-4B-Instruct-2507 on the same held out Banking77 rows at the same latency.
+    model_id: str = "Qwen/Qwen3-VL-4B-Instruct"
+    dtype: str = "bfloat16"
+    device: str = "cuda"
+    # None leaves the download location to HF_HOME, so installing this package puts
+    # nothing in whatever directory the caller happened to start in. A path keeps the
+    # weights beside the project instead, which is what the examples do.
+    models_dir: Path | None = None
+
+    served_model_id: str = "logit-classifier-0.1.0"
+    model_aliases: tuple[str, ...] = ("logit-latest", "logit-preview", "jev-latest", "jev-preview")
+
+    # None asks for the backend model's fitted value, resolved by Classifier.
+    temperature: float | None = None
+    use_prior_debias: bool = True
+    # None keeps the running prior in memory for the life of the process, so importing
+    # this library writes no file the host never asked for. from_env turns it back on.
+    calibration_path: Path | None = None
+
+    # "joint" scores all levels in one branch, "independent" judges each alone.
+    score_method: str = "joint"
+    # Averaging over several letterings cancels the model's preference for a label
+    # position. Measured on Banking77, four letterings moved accuracy from 0.554 to
+    # 0.693. Off by default because it multiplies the branch count.
+    permutations: int = 1
+    # An escape label absorbs the mass the model would otherwise spread over wrong
+    # options, so offering one raised accuracy from 0.881 to 0.887 as well as scoring
+    # how badly the model wants to decline.
+    abstain: bool = ABSTAIN
+
+    # False scores each branch in its own forward pass, which makes a question's
+    # logits independent of the sibling questions batched with it.
+    batch_branches: bool = True
+    max_batch_rows: int = 32
+
+    def __post_init__(self) -> None:
+        # A ComfyUI pack builds Config directly, so from_env cannot be the only check.
+        _reject_score_method("Config.score_method", self.score_method)
+        _reject_temperature("Config.temperature", self.temperature)
+        _reject_below_one("Config.permutations", self.permutations)
+        _reject_below_one("Config.max_batch_rows", self.max_batch_rows)
+
+    @classmethod
+    def from_env(cls) -> Config:
+        models_dir = os.environ.get("LOGIT_MODELS_DIR")
+        score_method = os.environ.get("LOGIT_SCORE_METHOD", cls.score_method)
+        permutations = _env_number("LOGIT_PERMUTATIONS", int)
+        temperature = _env_number("LOGIT_TEMPERATURE", float)
+
+        # Checked here as well as in __post_init__ so the message names the variable set.
+        _reject_score_method("LOGIT_SCORE_METHOD", score_method)
+        _reject_temperature("LOGIT_TEMPERATURE", temperature)
+        if permutations is not None:
+            _reject_below_one("LOGIT_PERMUTATIONS", permutations)
+        return cls(
+            model_id=os.environ.get("LOGIT_MODEL_ID", cls.model_id),
+            models_dir=Path(models_dir) if models_dir else None,
+            device=os.environ.get("LOGIT_DEVICE", cls.device),
+            temperature=temperature,
+            use_prior_debias=_env_switch("LOGIT_PRIOR_DEBIAS", cls.use_prior_debias),
+            batch_branches=_env_switch("LOGIT_BATCH_BRANCHES", cls.batch_branches),
+            score_method=score_method,
+            permutations=cls.permutations if permutations is None else permutations,
+            abstain=_env_switch("LOGIT_ABSTAIN", cls.abstain),
+            calibration_path=Path(os.environ.get("LOGIT_CALIBRATION_PATH", "calibration.json")),
+        )

logit_classifier/deps.py

@@ -0,0 +1,30 @@
+"""Loading the packages that live behind an extra.
+
+The core install names numpy alone so a ComfyUI pack never asks pip to resolve
+torch beside ComfyUI's pinned CUDA build. Everything outside that core is
+imported here, on first use, and a missing one fails with its install line.
+"""
+
+from __future__ import annotations
+
+import importlib
+from types import ModuleType
+
+from .errors import LogitClassifierError
+
+DISTRIBUTION = "logit-classifier"
+
+
+class MissingDependencyError(LogitClassifierError, ImportError):
+    """A path needs an optional extra that is not installed."""
+
+
+def require(module: str, extra: str) -> ModuleType:
+    """Import a module from an extra, or fail naming the command that installs it."""
+    try:
+        return importlib.import_module(module)
+    except ImportError as error:
+        raise MissingDependencyError(
+            f"{module} is needed here and is not installed. "
+            f'Install it with: pip install "{DISTRIBUTION}[{extra}]"'
+        ) from error

logit_classifier/errors.py

@@ -0,0 +1,16 @@
+"""The one base every error this library raises inherits.
+
+A host catching one name covers the library. An error the caller can fix by
+changing its input also inherits ValueError. An internal invariant failure also
+inherits RuntimeError. A missing optional dependency also inherits ImportError.
+"""
+
+from __future__ import annotations
+
+
+class LogitClassifierError(Exception):
+    """Any failure this library raises on purpose."""
+
+
+class ConfigError(LogitClassifierError, ValueError):
+    """A Config field or a LOGIT_ environment variable holds a value it cannot take."""

logit_classifier/labels.py

@@ -0,0 +1,91 @@
+"""Option label alphabet, boundary verification and the >52-option split."""
+
+from __future__ import annotations
+
+import string
+from dataclasses import dataclass
+from typing import Protocol
+
+from .errors import LogitClassifierError
+
+# Bare single letters are one token each in the Qwen vocabulary when the prefill
+# ends on "Ġ(", which gives 52 usable labels. Verified per tokenizer at startup.
+LABEL_ALPHABET = string.ascii_uppercase + string.ascii_lowercase
+MAX_LABELS_PER_BRANCH = len(LABEL_ALPHABET)
+
+
+class LabelBoundaryError(LogitClassifierError, RuntimeError):
+    """A label does not occupy exactly one token after the prefill."""
+
+
+class Encoder(Protocol):
+    """The one method the boundary proof needs, so it binds to no host."""
+
+    def encode(self, text: str) -> list[int]: ...
+
+
+def verify_label_ids(tokenizer: Encoder, prefill_text: str, count: int) -> list[int]:
+    """Return one token id per label, proving each adds exactly one token.
+
+    The check must run against the rendered prefill, not the label alone. A
+    prefill ending in a space silently merges the space into the letter, which
+    an isolated encode of "A" cannot detect.
+    """
+    base_ids: list[int] = []
+    label_ids: list[int] = []
+
+    if count > MAX_LABELS_PER_BRANCH:
+        raise LabelBoundaryError(f"{count} labels exceeds the {MAX_LABELS_PER_BRANCH} available")
+
+    base_ids = tokenizer.encode(prefill_text)
+    for label in LABEL_ALPHABET[:count]:
+        full = tokenizer.encode(prefill_text + label)
+        if len(full) != len(base_ids) + 1 or full[: len(base_ids)] != base_ids:
+            raise LabelBoundaryError(
+                f"label {label!r} does not add exactly one token after {prefill_text!r}"
+            )
+        label_ids.append(full[-1])
+
+    if len(set(label_ids)) != len(label_ids):
+        raise LabelBoundaryError("labels collide onto the same token id")
+    return label_ids
+
+
+@dataclass(frozen=True)
+class BranchPlan:
+    """How one choice question's options map onto scored branches.
+
+    A flat plan scores every option in a single branch. A split plan scores each
+    option group in its own branch with an escape label, and combine_escape weighs the
+    groups by the mass each keeps off that label.
+    """
+
+    groups: tuple[tuple[int, ...], ...]
+    split: bool
+
+
+def plan_branches(option_count: int, *, abstain: bool = False) -> BranchPlan:
+    """Split options into groups only when they exceed one branch's labels.
+
+    A split group holds one label back for the escape label that lets it say the
+    answer is not among its members. abstain means a flat plan carries that label
+    as well, so it has room for one option fewer.
+    """
+    group_count = 0
+    group_size = 0
+    capacity = MAX_LABELS_PER_BRANCH - 1
+    groups: list[tuple[int, ...]] = []
+
+    if option_count + (1 if abstain else 0) <= MAX_LABELS_PER_BRANCH:
+        return BranchPlan(groups=(tuple(range(option_count)),), split=False)
+
+    group_count = -(-option_count // capacity)
+    if group_count > capacity:
+        raise LabelBoundaryError(
+            f"{option_count} options needs {group_count} groups, over the "
+            f"{capacity} label limit"
+        )
+    group_size = -(-option_count // group_count)
+    for start in range(0, option_count, group_size):
+        groups.append(tuple(range(start, min(start + group_size, option_count))))
+    return BranchPlan(groups=tuple(groups), split=True)

logit_classifier/prompt.py

@@ -0,0 +1,154 @@
+"""Prompt assembly: one shared prefix, one suffix per scored branch.
+
+Everything here is text. Applying the host's chat template around it is the
+backend's job, since the HF tokenizer and a ComfyUI CLIP render it differently.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from .labels import LABEL_ALPHABET, BranchPlan, plan_branches
+from .schema import (
+    ChoiceQuestion,
+    NoulQuestion,
+    Question,
+    ScoreQuestion,
+    render_content,
+)
+
+# Bumped whenever a change to branch layout or prior bucketing makes an older learned
+# prior invalid.
+PROMPT_VERSION = 2
+
+SYSTEM_PROMPT = (
+    "You are a classification engine. Read the context, then answer the single "
+    "question about it.\n"
+    "Reply with exactly one option label from the list. Labels are case-sensitive."
+)
+
+# The processor replaces this single token with one per patch, so the marker costs
+# nothing to write and expands to the right length on its own.
+IMAGE_MARKER = "<|vision_start|><|image_pad|><|vision_end|>"
+
+# A split group carries this extra label so it can decline, which is what lets
+# groups be weighed against each other without a separate group branch.
+ESCAPE_LABEL = "none of these"
+ESCAPE_DESCRIPTION = "the answer is not in this list"
+
+BranchKind = Literal["choice", "member", "score_level", "score_joint", "noul"]
+
+
+@dataclass(frozen=True)
+class Branch:
+    """One prompt whose final token position carries a distribution."""
+
+    question_id: str
+    kind: BranchKind
+    label_count: int
+    # Which option or level indices this branch's labels map back to.
+    targets: tuple[int, ...]
+    suffix_text: str
+
+
+def _option_lines(entries: list[tuple[str, str]]) -> str:
+    lines = []
+    for index, (name, description) in enumerate(entries):
+        label = LABEL_ALPHABET[index]
+        lines.append(f"({label}) {name}: {description}" if description else f"({label}) {name}")
+    return "\n".join(lines)
+
+
+def _question_block(prompt_line: str, entries: list[tuple[str, str]]) -> str:
+    return f"{prompt_line}\nOptions:\n{_option_lines(entries)}"
+
+
+def _choice_branches(qid: str, question: ChoiceQuestion, plan: BranchPlan,
+                     abstain: bool = False) -> list[Branch]:
+    names: list[str] = list(question.criteria)
+    prompt_line = f"Question: {render_content(question.instructions)}".rstrip()
+    branches: list[Branch] = []
+
+    if not plan.split:
+        entries = [(n, render_content(question.criteria[n])) for n in names]
+        if not abstain:
+            return [Branch(qid, "choice", len(names), tuple(range(len(names))),
+                           _question_block(prompt_line, entries))]
+        # One group carrying the escape label, so the split path reads it the same way.
+        entries.append((ESCAPE_LABEL, ESCAPE_DESCRIPTION))
+        return [Branch(qid, "member", len(entries), tuple(range(len(names))),
+                       _question_block(prompt_line, entries))]
+
+    for group in plan.groups:
+        entries = [(names[i], render_content(question.criteria[names[i]])) for i in group]
+        entries.append((ESCAPE_LABEL, ESCAPE_DESCRIPTION))
+        branches.append(Branch(qid, "member", len(entries), tuple(group),
+                               _question_block(prompt_line, entries)))
+    return branches
+
+
+def _score_branches(qid: str, question: ScoreQuestion, method: str) -> list[Branch]:
+    """Expand a score question into branches.
+
+    "independent" follows Jev's stated semantics, judging each level on its own.
+    "joint" offers the levels as one labelled list, which forces the model to
+    discriminate between adjacent levels instead of accepting several at once.
+    Levels carry letters rather than indices either way, so no level number
+    reaches the model.
+    """
+    instructions = render_content(question.instructions)
+    branches: list[Branch] = []
+
+    if method == "joint":
+        entries = [(render_content(level), "") for level in question.criteria]
+        prompt_line = (
+            f"Question: {instructions}\n" if instructions else ""
+        ) + "Which description best matches the context?"
+        return [Branch(qid, "score_joint", len(question.criteria),
+                       tuple(range(len(question.criteria))),
+                       _question_block(prompt_line, entries))]
+
+    for index, level in enumerate(question.criteria):
+        prompt_line = (
+            f"Question: {instructions}\n" if instructions else ""
+        ) + f"Statement: the context matches this description: {render_content(level)}"
+        branches.append(Branch(qid, "score_level", 2, (index,),
+                               _question_block(prompt_line, [("yes", ""), ("no", "")])))
+    return branches
+
+
+def _noul_branches(qid: str, question: NoulQuestion) -> list[Branch]:
+    prompt_line = f"Statement: {render_content(question.instructions)}".rstrip()
+    yes_text = render_content(question.criteria.true_) if question.criteria else ""
+    no_text = render_content(question.criteria.false_) if question.criteria else ""
+    return [Branch(qid, "noul", 2, (0,),
+                   _question_block(prompt_line, [("yes", yes_text), ("no", no_text)]))]
+
+
+def build_branches(questions: dict[str, Question], score_method: str = "joint",
+                   abstain: bool = False) -> list[Branch]:
+    """Expand every question into the branches that must be scored."""
+    branches: list[Branch] = []
+
+    for qid, question in questions.items():
+        if isinstance(question, ChoiceQuestion):
+            branches.extend(_choice_branches(
+                qid, question, plan_branches(len(question.criteria), abstain=abstain), abstain
+            ))
+        elif isinstance(question, ScoreQuestion):
+            branches.extend(_score_branches(qid, question, score_method))
+        else:
+            branches.extend(_noul_branches(qid, question))
+    return branches
+
+
+def prefix_content(state: Any, has_image: bool = False) -> str:
+    """Build the user-message body every branch shares, image marker included."""
+    marker = f"{IMAGE_MARKER}\n" if has_image else ""
+    return f"Context:\n{marker}{render_content(state)}\n\n"
+
+
+def branch_content(state: Any, branch: Branch, has_image: bool = False) -> str:
+    """One branch's full user-message body, opening with the shared prefix."""
+    return prefix_content(state, has_image) + branch.suffix_text