add primitives.py

primitives.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python3
+"""
+primitives.py — the one place a row becomes a prompt.
+
+TypeSafe serves three question types against a state:
+
+  choice   pick one of several named options   -> probabilities + confidence
+  score    rate against ordered levels          -> expected level + legend
+  noul     is this true?                        -> a single probability
+
+All three arrive with `criteria`: a description per option, or per level. The
+student is asked to read those descriptions at serve time, so it has to be
+trained on them — and the teacher has to judge against the same text, or the
+soft label describes a prompt nobody will ever send. That is why build_data.py,
+teacher.py and jev_lite.py all render through this module and none of them
+format options themselves.
+
+Row shape (a superset of what jev_lite trains on):
+
+  {"type": "choice", "state": ..., "question": ...,
+   "options": ["billing", "technical"],
+   "criteria": {"billing": "Payments, invoicing, refunds", ...},
+   "label": [0.7, 0.3]}
+
+  {"type": "score", "ordered": true, "options": ["Calm", "Frustrated", "Angry"],
+   "label": [...]}                  # options ARE the levels; legend is positional
+
+  {"type": "noul", "options": ["true", "false"],
+   "criteria": {"true": "...", "false": "..."}, "label": [p_yes, p_no]}
+
+`criteria` is always optional: the API allows a bare option set, and some real
+tasks have no meaningful description to give (RACE's options are the answer
+text itself). Training on a mix teaches the model to use criteria when they are
+there and cope when they are not.
+"""
+
+LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+
+# noul answers in TypeSafe are keyed true/false, and the reported probability is
+# the one for "true" — so index 0 must be the affirmative, always.
+TRUE_FALSE = ["true", "false"]
+YES_WORDS = {"yes", "true", "y"}
+NO_WORDS = {"no", "false", "n"}
+
+PREFIX = "Read the state and answer the question with one option letter.\n\n<state>\n"
+
+
+def kind_of(row):
+    """choice | score | noul, from an explicit type or the shape of the row."""
+    if row.get("type") in ("choice", "score", "noul"):
+        return row["type"]
+    if row.get("ordered"):
+        return "score"
+    opts = [str(o).strip().lower() for o in row["options"]]
+    if len(opts) == 2 and any(o in YES_WORDS for o in opts) \
+            and any(o in NO_WORDS for o in opts):
+        return "noul"
+    return "choice"
+
+
+def normalize(row):
+    """Stamp `type`, and put noul rows in true/false order with true first.
+
+    Returns a new row; the answer index and soft label are permuted with the
+    options, so a normalized row means exactly what the original did.
+    """
+    row = dict(row)
+    kind = kind_of(row)
+    row["type"] = kind
+    if kind == "score":
+        row["ordered"] = True
+        return row
+    if kind != "noul":
+        return row
+
+    opts = [str(o).strip().lower() for o in row["options"]]
+    yes_at = next((i for i, o in enumerate(opts) if o in YES_WORDS), None)
+    no_at = next((i for i, o in enumerate(opts) if o in NO_WORDS), None)
+    if yes_at is None or no_at is None or yes_at == no_at:
+        row["type"] = "choice"          # not actually a yes/no pair
+        return row
+
+    perm = [yes_at, no_at]
+    criteria = row.get("criteria")
+    if isinstance(criteria, dict) and criteria:
+        row["criteria"] = {TRUE_FALSE[j]: criteria.get(row["options"][i])
+                           for j, i in enumerate(perm)
+                           if criteria.get(row["options"][i])}
+    row["options"] = list(TRUE_FALSE)
+    if "label" in row:
+        row["label"] = [row["label"][i] for i in perm]
+    if "answer" in row:
+        row["answer"] = perm.index(row["answer"])
+    return row
+
+
+def criterion_for(row, option, position):
+    """The description of one option: dict lookup for choice/noul, positional for score."""
+    criteria = row.get("criteria")
+    if not criteria:
+        return None
+    if isinstance(criteria, dict):
+        return criteria.get(option)
+    if isinstance(criteria, list) and position < len(criteria):
+        # A score's levels and its criteria are the same ordered list; only
+        # index into it when the row kept them apart.
+        return criteria[position] if criteria[position] != option else None
+    return None
+
+
+def option_lines(row, options=None):
+    """A. billing — Payments, invoicing, refunds"""
+    options = row["options"] if options is None else options
+    if len(options) > len(LETTERS):
+        raise ValueError(f"max {len(LETTERS)} options, got {len(options)}")
+    lines = []
+    for i, option in enumerate(options):
+        desc = criterion_for(row, option, i)
+        lines.append(f"{LETTERS[i]}. {option}" + (f" — {desc}" if desc else ""))
+    return "\n".join(lines)
+
+
+def question_block(row, options=None):
+    """Everything after the state: the question, its options, their criteria.
+
+    The header tells the model which primitive it is looking at. A score is the
+    one case where option ORDER carries meaning, so it says so out loud.
+    """
+    header = ("Levels, lowest to highest:" if kind_of(row) == "score"
+              else "Options:")
+    return f"Question: {row['question']}\n{header}\n{option_lines(row, options)}"
+
+
+def build_prompt(row, options=None):
+    """The full teacher-facing prompt for one row."""
+    return (f"{PREFIX}{row['state']}\n</state>\n\n{question_block(row, options)}\n\n"
+            "Reply with exactly one option letter and nothing else.\nAnswer:")
+
+
+# ------------------------------------------------------------------ answers
+
+def confidence(kind, probs, score=None):
+    """The model's probability that the answer it just gave is the right one.
+
+    One rule, realized per type, because each type returns a different thing:
+    a choice returns its argmax, so confidence is that option's probability; a
+    score returns an expected level, so it is the mass that rounds to the level
+    reported. Both are directly readable — 0.8 means wrong one time in five —
+    which is what makes a threshold an error budget instead of a vibe.
+
+    Chosen by measurement, not taste: over 1898 held-out rows these beat
+    entropy, margin, chance-corrected top and collision entropy on both AUROC
+    (ranking right answers above wrong ones) and ECE (the number meaning what
+    it says). Normalized entropy, the obvious first guess, was the worst of the
+    lot — it is dominated by small probabilities, so it reads a decisive
+    0.85/0.08/0.07 as barely-there confidence and sends good answers to review.
+    """
+    if kind == "score":
+        return sum(p for i, p in enumerate(probs) if abs(i - score) <= 0.5)
+    return max(probs)
+
+
+def answer(row, probs):
+    """Shape a probability distribution into the API's answer for this type."""
+    kind = kind_of(row)
+    options = row["options"]
+    # A noul carries no confidence field: the probability IS the answer, and
+    # the caller thresholds it directly.
+    if kind == "noul":
+        return {"type": "noul", "noul": round(probs[0], 4)}
+
+    if kind == "score":
+        expected = sum(i * p for i, p in enumerate(probs))
+        return {"type": "score",
+                "score": round(expected, 4),
+                "legend": {str(i): o for i, o in enumerate(options)},
+                "probabilities": {str(i): round(p, 4) for i, p in enumerate(probs)},
+                "confidence": round(confidence(kind, probs, expected), 4)}
+    return {"type": "choice",
+            "choice": options[max(range(len(probs)), key=probs.__getitem__)],
+            "probabilities": {o: round(p, 4) for o, p in zip(options, probs)},
+            "confidence": round(confidence(kind, probs), 4)}
+
+
+def temper(probs, temperature):
+    """Flatten (T>1) or sharpen (T<1) a distribution, leaving the argmax alone.
+
+    Serving-time calibration: an adapter trained against one numeric precision
+    and served at another produces a distribution of the wrong sharpness. The
+    ranking is unaffected, so accuracy does not move — only how confident the
+    answer claims to be, which is the part that gates actions.
+    """
+    import math
+    if temperature == 1.0:
+        return list(probs)
+    logs = [math.log(max(p, 1e-12)) / temperature for p in probs]
+    top = max(logs)
+    exp = [math.exp(x - top) for x in logs]
+    total = sum(exp)
+    return [x / total for x in exp]
+
+
+def normalized_entropy(probs):
+    """0 = certain, 1 = uniform. Divided by log(n) so option counts compare."""
+    import math
+    n = len(probs)
+    if n < 2:
+        return 0.0
+    h = -sum(p * math.log(p) for p in probs if p > 1e-12)
+    return h / math.log(n)