| """Map Gemma letter logits onto TypeSafe's three question primitives. |
| |
| The letter distribution is the model's raw output. Each primitive renormalizes |
| the distribution over its legal answer letters, so the probabilities sum to 1 |
| across the caller's options. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass, field |
|
|
| from .letters import LETTER_INDEX, LETTERS, _softmax |
|
|
| |
| |
| GENERIC_EXAMPLES: tuple[tuple[str, str, str], ...] = ( |
| ("I need this fixed immediately, it is blocking everything.", "Does the text express urgency?", "Y"), |
| ("Here is the documentation link for reference.", "Does the text express urgency?", "N"), |
| ("The app crashes when I open the settings page.", "Does the text contain a bug report?", "Y"), |
| ("Just wanted to share the meeting notes from today.", "Does the text contain a bug report?", "N"), |
| ("Please send this to a human agent right now.", "Does the text ask for a human agent?", "Y"), |
| ("No rush on this, whenever you get a chance.", "Does the text ask for a human agent?", "N"), |
| ("URGENT: the payment system is failing right now.", "Does the text express urgency?", "Y"), |
| ("The sky is blue and the grass is green.", "Does the text express urgency?", "N"), |
| ) |
|
|
| DEFAULT_EXAMPLE_COUNT = 6 |
|
|
|
|
| @dataclass(frozen=True) |
| class QuestionOption: |
| letter: str |
| label: str |
| description: str | None = None |
|
|
|
|
| @dataclass |
| class Question: |
| id: str |
| type: str |
| instructions: str |
| options: list[QuestionOption] = field(default_factory=list) |
| criteria: object = None |
|
|
|
|
| @dataclass(frozen=True) |
| class QuestionAnswer: |
| type: str |
| choice: str | None = None |
| probabilities: dict[str, float] = field(default_factory=dict) |
| confidence: float | None = None |
| noul: float | None = None |
| score: float | None = None |
| legend: dict[str, str] | None = None |
|
|
| def to_dict(self) -> dict: |
| payload: dict = {"type": self.type} |
| if self.type == "noul": |
| payload["noul"] = round(self.noul or 0.0, 6) |
| elif self.type == "choice": |
| payload["choice"] = self.choice |
| payload["probabilities"] = {k: round(v, 6) for k, v in self.probabilities.items()} |
| payload["confidence"] = round(self.confidence or 0.0, 6) |
| elif self.type == "score": |
| payload["score"] = round(self.score or 0.0, 6) |
| payload["legend"] = self.legend or {} |
| payload["probabilities"] = {k: round(v, 6) for k, v in self.probabilities.items()} |
| payload["confidence"] = round(self.confidence or 0.0, 6) |
| return payload |
|
|
|
|
| def confidence_from_probabilities(probabilities: list[float]) -> float: |
| """1.0 when all mass is on one option, 0.0 when uniform.""" |
| n = len(probabilities) |
| if n <= 1: |
| return 1.0 |
| entropy = -sum(p * math.log(p) for p in probabilities if p > 0) |
| return max(0.0, 1.0 - entropy / math.log(n)) |
|
|
|
|
| def _predicate_forms(instructions: str) -> tuple[str, str]: |
| """Turn an instruction into affirmative/negative option text. |
| |
| "Does the text request a refund?" -> ("Yes, the text requests a refund.", |
| "No, the text does not request a refund.") |
| """ |
| text = instructions.strip().rstrip(".").strip() |
| lowered = text.lower() |
|
|
| if lowered.startswith("does the text "): |
| body = text[len("Does the text ") :] |
| affirmative = f"the text {body}" |
| elif lowered.startswith("is the text "): |
| body = text[len("Is the text ") :] |
| affirmative = f"the text {body}" |
| elif lowered.startswith("is this "): |
| body = text[len("Is this ") :] |
| affirmative = f"this {body}" |
| elif lowered.startswith("does this "): |
| body = text[len("Does this ") :] |
| affirmative = f"this {body}" |
| elif lowered.startswith("does the "): |
| body = text[len("Does the ") :] |
| affirmative = f"the {body}" |
| elif lowered.startswith("is the "): |
| body = text[len("Is the ") :] |
| affirmative = f"the {body}" |
| elif lowered.startswith("did the "): |
| body = text[len("Did the ") :] |
| affirmative = f"the {body}" |
| elif lowered.startswith("can the "): |
| body = text[len("Can the ") :] |
| affirmative = f"the {body}" |
| else: |
| affirmative = text |
|
|
| affirmative = affirmative.rstrip("?").strip() |
| if affirmative.lower().startswith("the "): |
| negative = "the " + affirmative[4:].split(" ", 1)[0] + " does not" + ( |
| " " + affirmative[4:].split(" ", 1)[1] if " " in affirmative[4:] else "" |
| ) |
| else: |
| negative = f"it is not the case that {affirmative}" |
| return f"Yes, {affirmative}.", f"No, {negative}." |
|
|
|
|
| def build_question(qid: str, spec: dict) -> Question: |
| """Parse one TypeSafe question spec into a letter-addressable Question.""" |
| qtype = spec.get("type") |
| if qtype not in {"noul", "choice", "score"}: |
| raise ValueError(f"question {qid!r} has unsupported type {qtype!r}") |
|
|
| instructions = spec.get("instructions") |
| if not isinstance(instructions, str) or not instructions.strip(): |
| raise ValueError(f"question {qid!r} requires non-empty string instructions") |
|
|
| criteria = spec.get("criteria") |
|
|
| if qtype == "noul": |
| if isinstance(criteria, dict): |
| yes_desc = criteria.get("true") |
| no_desc = criteria.get("false") |
| else: |
| yes_desc = no_desc = None |
| yes_text, no_text = _predicate_forms(instructions) |
| return Question( |
| id=qid, |
| type=qtype, |
| instructions=instructions, |
| options=[ |
| QuestionOption("Y", "yes", yes_desc or yes_text), |
| QuestionOption("N", "no", no_desc or no_text), |
| ], |
| criteria=criteria, |
| ) |
|
|
| if qtype == "choice": |
| if not isinstance(criteria, dict) or not criteria: |
| raise ValueError(f"choice question {qid!r} requires a non-empty criteria map") |
| if len(criteria) > len(LETTERS): |
| raise ValueError(f"choice question {qid!r} supports at most 26 options") |
| entries = list(criteria.items()) |
| else: |
| if not isinstance(criteria, list) or len(criteria) < 2: |
| raise ValueError(f"score question {qid!r} requires an ordered list of at least two levels") |
| if len(criteria) > len(LETTERS): |
| raise ValueError(f"score question {qid!r} supports at most 26 levels") |
| entries = [(f"level {i}", level) for i, level in enumerate(criteria)] |
|
|
| options = [ |
| QuestionOption( |
| LETTERS[i], |
| label, |
| description if isinstance(description, str) else str(description), |
| ) |
| for i, (label, description) in enumerate(entries) |
| ] |
| return Question(id=qid, type=qtype, instructions=instructions, options=options, criteria=criteria) |
|
|
|
|
| def score_question(question: Question, letter_logits: dict[str, float], temperature: float = 1.0) -> QuestionAnswer: |
| """Turn raw letter logits into a typed answer, renormalized over legal options.""" |
| option_letters = [option.letter for option in question.options] |
| missing = [letter for letter in option_letters if letter not in letter_logits] |
| if missing: |
| raise ValueError(f"question {question.id!r} is missing logits for {missing}") |
|
|
| logits = [letter_logits[letter] for letter in option_letters] |
| probs = _softmax(logits, temperature) |
|
|
| if question.type == "noul": |
| noul = probs[0] |
| return QuestionAnswer( |
| type="noul", |
| noul=noul, |
| ) |
|
|
| if question.type == "choice": |
| pairs = sorted(zip(question.options, probs), key=lambda op: op[1], reverse=True) |
| top = pairs[0] |
| return QuestionAnswer( |
| type="choice", |
| choice=top[0].label, |
| probabilities={option.label: p for option, p in pairs}, |
| confidence=confidence_from_probabilities(probs), |
| ) |
|
|
| score = sum(i * p for i, p in enumerate(probs)) |
| return QuestionAnswer( |
| type="score", |
| score=score, |
| probabilities={str(i): p for i, p in enumerate(probs)}, |
| legend={str(i): option.description or "" for i, option in enumerate(question.options)}, |
| confidence=confidence_from_probabilities(probs), |
| ) |
|
|
|
|
| def score_question_tensor(question: Question, letter_logits, temperature: float = 1.0) -> QuestionAnswer: |
| """Score a ``[26]`` device tensor, transferring only final answer values. |
| |
| The logits, legal-option gather, softmax, entropy, ranking, and score |
| reduction all remain on the tensor's device. Constructing ``QuestionAnswer`` |
| necessarily creates Python/JSON values at the API boundary. |
| """ |
| if temperature <= 0: |
| raise ValueError("temperature must be > 0") |
| if letter_logits.dim() != 1 or letter_logits.shape[0] != len(LETTERS): |
| raise ValueError("expected one [26] letter-logit tensor") |
|
|
| import torch |
|
|
| indices = torch.tensor( |
| [LETTER_INDEX[option.letter] for option in question.options], |
| device=letter_logits.device, |
| dtype=torch.long, |
| ) |
| probs = torch.softmax(letter_logits.index_select(0, indices).float() / temperature, dim=0) |
|
|
| if question.type == "noul": |
| return QuestionAnswer(type="noul", noul=float(probs[0].item())) |
|
|
| if len(question.options) <= 1: |
| confidence_value = 1.0 |
| else: |
| entropy = -(probs * torch.where(probs > 0, probs.log(), torch.zeros_like(probs))).sum() |
| confidence = 1.0 - entropy / torch.log(torch.tensor(float(len(question.options)), device=probs.device)) |
| confidence_value = float(torch.clamp(confidence, min=0.0).item()) |
|
|
| if question.type == "choice": |
| |
| order = torch.argsort(probs, descending=True, stable=True).tolist() |
| probability_values = probs.tolist() |
| return QuestionAnswer( |
| type="choice", |
| choice=question.options[order[0]].label, |
| probabilities={question.options[index].label: probability_values[index] for index in order}, |
| confidence=confidence_value, |
| ) |
|
|
| probability_values = probs.tolist() |
| levels = torch.arange(len(question.options), device=probs.device, dtype=probs.dtype) |
| score = float((levels * probs).sum().item()) |
| return QuestionAnswer( |
| type="score", |
| score=score, |
| probabilities={str(index): value for index, value in enumerate(probability_values)}, |
| legend={str(index): option.description or "" for index, option in enumerate(question.options)}, |
| confidence=confidence_value, |
| ) |
|
|
