src/solomon/calibration.py

"""The readout calibration: one positive scalar per answer type, and nothing else.

Split out of `solomon/service.py` unchanged. A temperature rescales a probability; it cannot reorder one,
so nothing here can decline an answer. `ReadoutTemperature` refuses an artifact that carries anything with
structure beyond the five scalars and the head_key map, which is what a smuggled decision rule would look
like.
"""
import copy
import hashlib
import json
import math
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]

# Per-task readout temperature. Not a gate and not a decision: one positive scalar per task, applied to the logits
# before the softmax, which is the only parameter that fixes the MAGNITUDE of the returned probability. It cannot
# change any ranking (a positive temperature is monotone), so it cannot reintroduce abstention by another name.
TEMPERATURE_SCHEMA = 'solomon-readout-temperature-v2'
# v1 is accepted structurally so an older artifact still LOADS rather than crashing; which fit may actually SHIP
# is decided by the release tooling:RULED_TEMPERATURES, on values rather than on a version string.
TEMPERATURE_SCHEMAS = ('solomon-readout-temperature-v3', 'solomon-readout-temperature-v2', 'solomon-readout-temperature-v1')
TEMPERATURE = 1.0                       # served when no calibration artifact is bound
TEMPERATURE_TASKS = ('boolean', 'entity', 'multilabel', 'single', 'ordered')
# v1.1: the answer types the public API serves. A v3 temperature artifact must cover exactly these; it MAY also carry
# an entity scalar (the fit still reports the type), which is recorded and never served.
SERVED_TASKS = ('boolean', 'multilabel', 'single', 'ordered')
def digest(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(',', ':'), allow_nan=False).encode()).hexdigest()


class ReadoutTemperature:
    """A frozen `solomon-readout-temperature-v1`: one positive scalar per task, nothing else.

    This is a fitted artifact, and that is deliberate. The object removed from this path today was a
    5,125-parameter logistic regression over the full hidden state fitted on ~1 effective negative, whose
    regularisation was chosen on a panel containing zero errors; it separated perfectly and did not replicate.
    This is five scalars fitted on thousands of units, stable across the two disjoint fit panels, and it decides
    nothing -- it rescales a number. The risk classes are not comparable. It is still bound by the runtime
    identity like everything else, so it cannot drift from the stack it was measured on.

    Application (must match abstention_refit/readout.py, which is what the scalars were fitted on):
      Noul   -> p_yes(t) = sigmoid(z/t) on the COLLAPSED binary logit z (solomon.semantics.p_yes)
      Choice -> softmax(logits[:n]/t) over the listed options only (solomon.semantics.listed_probs)
    """
    # Nothing that could gate, rank-order differently, or make a certified claim may ride along with the scalars.
    # Prose is fine: the artifact documents itself, and `application.noul.confidence` is a sentence describing
    # what a caller reads off the tempered probability, not a fitted object. So a refused name is only refused
    # when it carries STRUCTURE (numbers, arrays, objects) -- that is what a smuggled head would look like.
    REFUSED = ('thresholds', 'threshold', 'weights', 'mean', 'scale', 'confidence', 'correctness',
               'feature_recipe', 'abstention', 'coverage', 'error_upper')
    REFUSED_ANYWHERE = ('weights', 'thresholds', 'correctness', 'abstention')

    def __init__(self, artifact, file_sha256=None):
        if not isinstance(artifact, dict) or artifact.get('schema') not in TEMPERATURE_SCHEMAS:
            raise ValueError('one of ' + ', '.join(TEMPERATURE_SCHEMAS) + ' required')
        if artifact.get('frozen') is not True:
            raise ValueError('readout temperature artifact must be frozen')
        stated = artifact.get('sha256')
        if stated != hashlib.sha256(json.dumps({k: v for k, v in artifact.items() if k != 'sha256'},
                                               sort_keys=True, separators=(',', ':')).encode()).hexdigest():
            raise ValueError('readout temperature checksum mismatch')
        bad = sorted(_structural_keys(artifact) & set(self.REFUSED))
        if bad:
            raise ValueError('readout temperature artifact carries non-temperature machinery: ' + ', '.join(bad))
        anywhere = sorted(_all_keys(artifact) & set(self.REFUSED_ANYWHERE))
        if anywhere:
            raise ValueError('readout temperature artifact names forbidden machinery: ' + ', '.join(anywhere))
        if set(artifact) & set(self.REFUSED):
            raise ValueError('readout temperature artifact carries a refused top-level field')
        models = artifact.get('models') or {}
        if artifact['schema'] == TEMPERATURE_SCHEMAS[0]:
            # v3 (v1.1): exactly the served types, plus an optional (never served) entity scalar.
            if not set(SERVED_TASKS) <= set(models) or set(models) - set(TEMPERATURE_TASKS):
                raise ValueError('readout temperature v3 needs one model per served task: ' + ', '.join(sorted(SERVED_TASKS))
                                 + ' (entity optional)')
        elif set(models) != set(TEMPERATURE_TASKS):
            raise ValueError('readout temperature needs exactly one model per task: ' + ', '.join(sorted(TEMPERATURE_TASKS)))
        values = {}
        for task, model in models.items():
            t = model.get('temperature')
            if isinstance(t, bool) or not isinstance(t, (int, float)) or not math.isfinite(t) or not 0 < t <= 20:
                raise ValueError(f'{task}: temperature must be a finite scalar in (0, 20]')
            values[task] = float(t)
        self.artifact = copy.deepcopy(artifact)
        self.values = values
        # TWO DIFFERENT HASHES, and confusing them makes a sound provenance chain look tampered with:
        #   payload_sha256 -- the artifact's own `sha256` field, over its contents with that field removed. A
        #                     self-referential field cannot hash the file containing it, so this is what it covers,
        #                     and it is what this class verifies (above) and what survives being embedded in the
        #                     serving binding, where there is no file at all.
        #   file_sha256    -- sha256 of the .json file on disk. This is what `shasum -a 256 <file>` returns and
        #                     what a reader verifying the standalone artifact will compute. Known only when we
        #                     loaded from a path; recorded in the binding's provenance either way.
        self.payload_sha256 = stated
        self.file_sha256 = file_sha256
        # The artifact indexes by head_key, which is the field the runtime carries on every branch. The mapping to
        # the five tasks is one-to-one, so a served branch whose head_key is not in the map is a readout the
        # scalars were never fitted against (e.g. a different choice design) -- that must fail, not silently use 1.0.
        # v3: a head_key may serve several types (the merged yes/no head serves boolean and multilabel), each with its
        # own scalar; the map then names a list. v1/v2: one type per head_key.
        self.by_head_key = dict((artifact.get('application') or {}).get('by_head_key') or {})
        for key, task in self.by_head_key.items():
            for t in (task if isinstance(task, list) else [task]):
                if t not in values:
                    raise ValueError(f'readout temperature by_head_key maps {key} to unknown task {t}')

    def is_identity(self):
        """All temperatures exactly 1.0: a no-op, and therefore safe on any model."""
        return all(v == 1.0 for v in self.values.values())

    def temperature(self, task, head_keys=()):
        if task not in self.values:
            raise ValueError(f'readout temperature has no scalar for task {task}')
        for key in {k for k in head_keys if k}:
            if self.by_head_key and key not in self.by_head_key:
                raise ValueError(f'readout temperature was not fitted for head_key {key!r} (fitted: '
                                 + ', '.join(sorted(self.by_head_key)) + '); refusing to serve it uncalibrated')
            fitted = self.by_head_key.get(key)
            if self.by_head_key and task not in (fitted if isinstance(fitted, list) else [fitted]):
                raise ValueError(f'head_key {key!r} is calibrated for task {fitted}, served as {task}')
        return self.values[task]

    def describe(self):
        return {'schema': self.artifact['schema'], 'payload_sha256': self.payload_sha256,
                'file_sha256': self.file_sha256, 'per_task': dict(self.values),
                'by_head_key': dict(self.by_head_key),
                'applied': 'logits are divided by the task temperature before the softmax; Noul temperature '
                           'applies to the collapsed binary logit (p_yes = sigmoid(z/T)), choice temperature to '
                           'the listed options (softmax(logits[:n]/T)); applied exactly once per read'}


class NoTemperature:
    """No calibration artifact bound: the readout is served raw at T = 1."""
    payload_sha256 = None
    file_sha256 = None

    def is_identity(self):
        return True

    def temperature(self, task, head_keys=()):
        return TEMPERATURE

    def describe(self):
        return {'schema': None, 'payload_sha256': None, 'file_sha256': None, 'per_task': {}, 'by_head_key': {},
                'applied': 'none: raw readout at T = 1'}


def load_temperature(source):
    """None -> NoTemperature; a dict or a path -> ReadoutTemperature (invalid fails closed).

    Loading from a path also records the file hash, which is what `shasum` returns; loading from a dict (the
    embedded copy inside a serving binding) cannot know it, and the binding's provenance carries it instead.
    """
    if source is None:
        return NoTemperature()
    if hasattr(source, 'temperature'):
        return source
    if isinstance(source, dict):
        return ReadoutTemperature(source)
    path = Path(source)
    if not path.is_absolute():
        path = ROOT / path
    if not path.exists():
        raise ValueError(f'readout temperature artifact not found: {path}')
    raw = path.read_bytes()
    return ReadoutTemperature(json.loads(raw), file_sha256=hashlib.sha256(raw).hexdigest())


def _structural_keys(value):
    """Keys anywhere in `value` whose own value is not a descriptive string (i.e. carries structure)."""
    out = set()
    if isinstance(value, dict):
        for k, v in value.items():
            if not isinstance(v, str):
                out.add(k)
            out |= _structural_keys(v)
    elif isinstance(value, list):
        for v in value:
            out |= _structural_keys(v)
    return out


def _all_keys(value):
    if isinstance(value, dict):
        return set(value) | (set().union(*(_all_keys(v) for v in value.values())) if value else set())
    if isinstance(value, list):
        return set().union(*(_all_keys(v) for v in value)) if value else set()
    return set()