mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
Marius CiepluchandClaude Fable 5.1
fbc3220bbe
learn/experiments.py: YAML manifests under plane/experiments/<id>.yaml, digest-checked at register, run and analyse time; `experiment`, `arm_outcome` and `experiment_result` rows; paired arms in a seeded repository order with a per-unit coin; arm A replay-only at $0; canary replicates under a fresh cache salt; repository-cluster paired bootstrap, exact McNemar, two O'Brien-Fleming looks, adopt/reject/equivalent/ inconclusive. `ousast learn experiment register|units|run|analyse`. Program: `ProgramSpec.sampling = retrieval_ensemble` (k samples at temperature 0 over k disjoint retrieved example sets; default unchanged); `Caller.salt`; `retrieve(exclude=)`. verify: OUSAST_VERIFY_CALLERS=off. Manifests exp-001-known-callers (plane arms, units pending) and exp-002-retrieval-ensemble (slice-2 incumbents, ceiling $4) -- not run. 18 scripted tests; module audit regenerated (151 modules). Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
531 lines
24 KiB
Python
531 lines
24 KiB
Python
"""The decision engine's program: a declared language-model classifier with local memory (design 4.1, Req 3.1).
|
|
|
|
DSPy-style and in-house. The **signature** (:data:`SIGNATURE`) declares what goes in -- the candidate's code
|
|
excerpt, its allow-listed signals, its inferred roles, its family -- and what comes out: ``verdict``
|
|
(``vulnerable | not_vulnerable | unsure``), ``family``, ``confidence`` in [0, 1] and a ``rationale`` of at most three
|
|
sentences citing at least one line of the code. The **modules** are ``Retrieve -> Classify(k)``:
|
|
|
|
- the prompt is the instruction plus the compiled demonstrations -- a prefix that is byte-identical for every
|
|
candidate, so the provider's prefix cache hits -- then the retrieved examples, then the candidate;
|
|
- sample 1 runs at temperature 0, samples 2..k at 0.7; an answer that does not parse, or cites no line in range, is
|
|
asked once more at temperature 0 and otherwise becomes ``unsure`` with ``parse_failed`` -- never ``not_vulnerable``;
|
|
- the raw score is ``s = mean_j p_j`` (``p_j`` the confidence of a ``vulnerable`` answer, one minus it for
|
|
``not_vulnerable``, 0.5 for ``unsure``); the verdict is the majority (ties are ``unsure``), and a majority
|
|
``unsure`` can never BLOCK.
|
|
|
|
Signals reach the prompt only through :func:`.schema.validate_x` (no ``EXCLUDED_FIELDS``, no free text) and the input
|
|
profile (``ProgramSpec.inputs``, default ``v1``: the engine instrument is withheld from prompts and from the retrieval
|
|
distance, see :data:`.schema.INPUT_PROFILES`); a missing
|
|
instrument is shown in words. Every model call goes through :class:`Caller`: a response cache keyed by
|
|
``sha256(model, Model parameters digest, messages, temperature, sample, json_object)`` stored as
|
|
``responses/<sha>.json``, so a rerun on the same inputs replays at $0 and reproduces byte for byte. Before a call the
|
|
rendered example and demonstration ids are checked against the evaluation boundary
|
|
(:func:`.retrieve.assert_boundary`).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import re
|
|
from collections.abc import Callable, Collection, Mapping, Sequence
|
|
from dataclasses import asdict, dataclass, field, replace
|
|
from typing import Any
|
|
|
|
from ..plane.budget import cost_of, prices_from
|
|
from ..plane.memory import MemoryStore
|
|
from .examples import Example
|
|
from .folds import Fold
|
|
from .retrieve import Retrieval, Target, assert_boundary, demonstrations_for, retrieve
|
|
from .schema import DEFAULT_INPUTS, features_for, instruments_for, validate_x, withheld
|
|
|
|
VERDICTS = ("vulnerable", "not_vulnerable", "unsure")
|
|
UNSURE = "unsure"
|
|
TEMPERATURES = (0.0, 0.7)
|
|
# How the k samples differ: by temperature (the default: sample 1 at 0, samples 2..k at 0.7 over one retrieved
|
|
# example set), or by the retrieved examples (every sample at temperature 0, each over a disjoint example set --
|
|
# the variance-reduction arm of experiment exp-002; design section 6, engine-level experiments).
|
|
SAMPLING = ("temperature", "retrieval_ensemble")
|
|
RETRIEVAL_ENSEMBLE = "retrieval_ensemble"
|
|
MAX_OUTPUT_TOKENS = 300
|
|
CHARS_PER_TOKEN = 4.0
|
|
STATE_WORDS = {"none": "no coverage", "failed": "failed", "not_applicable": "not applicable"}
|
|
_LINE_NUMBER = re.compile(r"^\s*(\d+) ")
|
|
|
|
|
|
class ReplayMiss(RuntimeError):
|
|
"""A replay-only run (no client) met a request that is not in the response cache."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Field:
|
|
name: str
|
|
direction: str # in | out
|
|
type: str
|
|
bound: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Signature:
|
|
fields: tuple[Field, ...]
|
|
|
|
@property
|
|
def digest(self) -> str:
|
|
return hashlib.sha256(json.dumps([asdict(f) for f in self.fields], sort_keys=True).encode()).hexdigest()
|
|
|
|
|
|
SIGNATURE = Signature(
|
|
(
|
|
Field("code", "in", "excerpt", "numbered lines, <= 80 lines and 4,000 characters; a delta adds the diff (<= 40 lines)"),
|
|
Field(
|
|
"signals",
|
|
"in",
|
|
"record",
|
|
"the allow-listed features the input profile shows, as name: value lines; a missing instrument in words",
|
|
),
|
|
Field("roles", "in", "roles", "inferred roles of the function: role, operation, line, origin"),
|
|
Field("family", "in", "family", "the candidate's family, or unknown"),
|
|
Field("verdict", "out", "enum", "vulnerable | not_vulnerable | unsure"),
|
|
Field("family", "out", "family", "one of the families, or none"),
|
|
Field("confidence", "out", "float", "[0, 1]: the probability that the verdict is right"),
|
|
Field("rationale", "out", "text", "at most 3 sentences; cites at least one line number of code (cited_lines)"),
|
|
)
|
|
)
|
|
|
|
# The answer contract is part of the signature, not of the searched instruction: a proposed instruction that leaves it
|
|
# out still gets it (``render_prefix``), so every prompt asks for JSON (DeepSeek's JSON mode refuses a prompt that
|
|
# never says "json") and every answer can parse.
|
|
ANSWER_FORMAT = (
|
|
'Answer with one JSON object: {"verdict": "vulnerable" | "not_vulnerable" | "unsure", "family": <family or\n'
|
|
'"none">, "confidence": <probability in [0, 1] that your verdict is right>, "rationale": <at most 3 sentences>,\n'
|
|
'"cited_lines": [<line numbers of the code that support the verdict>]}. Cite at least one line. Say "unsure"\n'
|
|
"when the code does not decide it."
|
|
)
|
|
BASELINE_INSTRUCTION = (
|
|
"You decide whether a function contains an exploitable security vulnerability of the given family.\n"
|
|
"You see the function's code (numbered lines), the signals static instruments emitted for it, and the roles\n"
|
|
"inferred for its calls. Signals are evidence, not verdicts: rules over-report, and a missing instrument is not a\n"
|
|
"negative. Similar labelled functions from other repositories follow as examples; reason about this code.\n" + ANSWER_FORMAT
|
|
)
|
|
|
|
|
|
# --- rendering ---------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def render_signals(x: Mapping[str, Any], instruments: Mapping[str, Mapping[str, Any]], profile: str, inputs: str = DEFAULT_INPUTS) -> str:
|
|
"""``name: value`` per feature of an instrument that ran; ``<instrument>: no coverage | failed | not applicable``
|
|
otherwise. The record is validated first: an unknown, label or identity field never reaches a prompt. An
|
|
instrument the input profile withholds (:data:`.schema.INPUT_PROFILES`; v1: the engine) is not rendered at all --
|
|
neither its state nor its values -- though the stored record keeps it."""
|
|
validate_x(x, instruments, profile)
|
|
hidden = withheld(inputs)
|
|
lines: list[str] = []
|
|
for name in instruments_for(profile):
|
|
if name == "language" or name in hidden:
|
|
continue
|
|
state = str(instruments[name]["state"])
|
|
if state != "ran":
|
|
lines.append(f"{name}: {STATE_WORDS[state]}")
|
|
continue
|
|
for spec in features_for(profile):
|
|
if spec.instrument == name and spec.name not in hidden:
|
|
value = x[spec.name]
|
|
lines.append(f"{spec.name}: {'null' if value is None else json.dumps(value)}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def render_roles(roles: Sequence[Mapping[str, Any]]) -> str:
|
|
if not roles:
|
|
return "none inferred"
|
|
keys = ("role", "kind", "operation", "line", "origin")
|
|
return "\n".join(", ".join(f"{k}={r[k]}" for k in keys if k in r and k != "path") for r in roles)
|
|
|
|
|
|
def _case(
|
|
code: str,
|
|
x: Mapping[str, Any],
|
|
instruments: Mapping[str, Mapping[str, Any]],
|
|
roles: Sequence[Mapping[str, Any]],
|
|
family: str,
|
|
profile: str,
|
|
language: str,
|
|
inputs: str = DEFAULT_INPUTS,
|
|
) -> str:
|
|
signals = render_signals(x, instruments, profile, inputs)
|
|
return f"Family: {family}\nLanguage: {language}\nCode:\n{code.rstrip()}\nSignals:\n{signals}\nRoles:\n{render_roles(roles)}"
|
|
|
|
|
|
def label_word(label: int) -> str:
|
|
return "vulnerable" if label == 1 else "not_vulnerable"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Demo:
|
|
"""A compiled demonstration: an example of ``C_boot`` with the rationale the program gave without seeing its label."""
|
|
|
|
example_id: str
|
|
excerpt_sha: str
|
|
label: int
|
|
family: str
|
|
verdict: str
|
|
confidence: float
|
|
rationale: str
|
|
cited_lines: tuple[int, ...]
|
|
alternate: Demo | None = None
|
|
|
|
def as_dict(self) -> dict[str, Any]:
|
|
out = {k: v for k, v in asdict(self).items() if k != "alternate"}
|
|
out["cited_lines"] = list(self.cited_lines)
|
|
out["alternate"] = self.alternate.as_dict() if self.alternate is not None else None
|
|
return out
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Mapping[str, Any]) -> Demo:
|
|
alt = data.get("alternate")
|
|
return cls(
|
|
str(data["example_id"]), str(data["excerpt_sha"]), int(data["label"]), str(data["family"]), str(data["verdict"]),
|
|
float(data["confidence"]), str(data["rationale"]), tuple(int(n) for n in data.get("cited_lines") or ()),
|
|
cls.from_dict(alt) if isinstance(alt, Mapping) else None,
|
|
) # fmt: skip
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProgramSpec:
|
|
"""What a compiled program fixes: instruction, demonstrations, model and sampling, retrieval settings."""
|
|
|
|
instruction: str = BASELINE_INSTRUCTION
|
|
demos: tuple[Demo, ...] = ()
|
|
profile: str = "static"
|
|
model: str = "deepseek-flash"
|
|
k: int = 1
|
|
k_ret: int = 6
|
|
lam: float = 0.5
|
|
n1: int = 40
|
|
priors: str = "off"
|
|
inputs: str = DEFAULT_INPUTS # the input profile: which instruments the classifier never sees (v1: the engine)
|
|
max_output_tokens: int = MAX_OUTPUT_TOKENS
|
|
sampling: str = SAMPLING[0] # temperature | retrieval_ensemble (an experiment arm's setting; never compiled)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Candidate:
|
|
"""A candidate as the program sees it: retrieval's target and the rendered (bounded, redacted) excerpt."""
|
|
|
|
target: Target
|
|
code: str
|
|
language: str
|
|
roles: tuple[Mapping[str, Any], ...] = ()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Prompt:
|
|
messages: list[dict[str, object]]
|
|
example_ids: tuple[str, ...]
|
|
demo_ids: tuple[str, ...]
|
|
valid_lines: frozenset[int]
|
|
|
|
@property
|
|
def prefix(self) -> str:
|
|
return str(self.messages[0]["content"])
|
|
|
|
@property
|
|
def estimated_tokens(self) -> int:
|
|
return round(sum(len(str(m["content"])) for m in self.messages) / CHARS_PER_TOKEN)
|
|
|
|
|
|
def line_numbers(code: str) -> frozenset[int]:
|
|
"""The source line numbers a rendered excerpt shows (the diff part has none)."""
|
|
return frozenset(int(m.group(1)) for line in code.split("\ndiff:\n", 1)[0].splitlines() if (m := _LINE_NUMBER.match(line)))
|
|
|
|
|
|
ExcerptText = Callable[[str], str | None]
|
|
|
|
|
|
def render_prefix(spec: ProgramSpec, demos: Sequence[Demo], examples: Mapping[str, Example], excerpt_text: ExcerptText) -> str:
|
|
"""The instruction and the demonstrations: identical for every candidate of one program and fold."""
|
|
instruction = spec.instruction.strip()
|
|
parts = [instruction if ANSWER_FORMAT in instruction else f"{instruction}\n{ANSWER_FORMAT}"]
|
|
for number, demo in enumerate(demos, 1):
|
|
example = examples[demo.example_id]
|
|
code = excerpt_text(demo.excerpt_sha) or ""
|
|
answer = {"verdict": demo.verdict, "family": demo.family, "confidence": demo.confidence, "rationale": demo.rationale,
|
|
"cited_lines": list(demo.cited_lines)} # fmt: skip
|
|
parts.append(
|
|
f"Demonstration {number}:\n"
|
|
+ _case(code, example.x, example.instruments, example.roles, example.family, spec.profile, example.language, spec.inputs)
|
|
+ f"\nAnswer: {json.dumps(answer, sort_keys=True)}"
|
|
)
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
def render(
|
|
spec: ProgramSpec,
|
|
candidate: Candidate,
|
|
retrieved: Sequence[Example],
|
|
demos: Sequence[Demo],
|
|
examples: Mapping[str, Example],
|
|
excerpt_text: ExcerptText,
|
|
) -> Prompt:
|
|
"""The prompt: [system: instruction + demonstrations] [user: retrieved examples (with their labels) + the candidate
|
|
(without one)]."""
|
|
blocks = []
|
|
for number, example in enumerate(retrieved, 1):
|
|
code = excerpt_text(example.excerpt_sha) or ""
|
|
blocks.append(
|
|
f"Example {number} (label: {label_word(example.label)}):\n"
|
|
+ _case(code, example.x, example.instruments, example.roles, example.family, spec.profile, example.language, spec.inputs)
|
|
)
|
|
t = candidate.target
|
|
blocks.append(
|
|
"Candidate:\n" + _case(candidate.code, t.x, t.instruments, candidate.roles, t.family, spec.profile, candidate.language, spec.inputs)
|
|
)
|
|
messages: list[dict[str, object]] = [
|
|
{"role": "system", "content": render_prefix(spec, demos, examples, excerpt_text)},
|
|
{"role": "user", "content": "\n\n".join(blocks)},
|
|
]
|
|
return Prompt(messages, tuple(e.id for e in retrieved), tuple(d.example_id for d in demos), line_numbers(candidate.code))
|
|
|
|
|
|
# --- parsing -----------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Answer:
|
|
verdict: str
|
|
family: str
|
|
confidence: float
|
|
rationale: str
|
|
cited_lines: tuple[int, ...]
|
|
parse_failed: bool = False
|
|
|
|
@property
|
|
def p(self) -> float:
|
|
"""The answer's probability that the candidate is vulnerable."""
|
|
if self.verdict == "vulnerable":
|
|
return self.confidence
|
|
if self.verdict == "not_vulnerable":
|
|
return 1.0 - self.confidence
|
|
return 0.5
|
|
|
|
|
|
PARSE_FAILED = Answer(UNSURE, "none", 0.0, "", (), parse_failed=True)
|
|
|
|
|
|
def parse_answer(content: str | None, valid_lines: frozenset[int], families: Sequence[str] = ()) -> Answer | None:
|
|
"""The answer, or ``None`` when it does not parse: not one JSON object, a verdict outside the vocabulary, a
|
|
confidence outside [0, 1], no rationale, or no cited line inside the excerpt."""
|
|
from ..provider.openrouter import OpenRouterError, parse_json_content
|
|
|
|
try:
|
|
data = parse_json_content(content or "")
|
|
except OpenRouterError:
|
|
return None
|
|
if not isinstance(data, dict) or data.get("verdict") not in VERDICTS:
|
|
return None
|
|
confidence = data.get("confidence")
|
|
if isinstance(confidence, bool) or not isinstance(confidence, int | float) or not 0.0 <= float(confidence) <= 1.0:
|
|
return None
|
|
rationale = data.get("rationale")
|
|
if not isinstance(rationale, str) or not rationale.strip():
|
|
return None
|
|
cited = data.get("cited_lines")
|
|
lines = (
|
|
tuple(sorted({int(n) for n in cited if isinstance(n, int) and not isinstance(n, bool) and n in valid_lines}))
|
|
if isinstance(cited, list)
|
|
else ()
|
|
)
|
|
if not lines:
|
|
return None
|
|
family = data.get("family")
|
|
family = family if isinstance(family, str) and (not families or family in families or family == "none") else "none"
|
|
return Answer(str(data["verdict"]), family, round(float(confidence), 6), rationale.strip()[:600], lines)
|
|
|
|
|
|
# --- calls and the response cache --------------------------------------------------------------------------------------
|
|
|
|
|
|
def request_key(
|
|
model: str, params_digest: str, messages: Sequence[Mapping[str, object]], temperature: float, sample: str, json_object: bool
|
|
) -> str:
|
|
canonical = {"json_object": json_object, "messages": list(messages), "model": model, "params": params_digest, "sample": sample,
|
|
"temperature": temperature} # fmt: skip
|
|
return hashlib.sha256(json.dumps(canonical, sort_keys=True, ensure_ascii=False).encode("utf-8")).hexdigest()
|
|
|
|
|
|
def params_digest(parameters: Mapping[str, Any] | None) -> str:
|
|
return hashlib.sha256(json.dumps(dict(parameters or {}), sort_keys=True).encode()).hexdigest()
|
|
|
|
|
|
@dataclass
|
|
class Caller:
|
|
"""Every model call of a program run: cached by request in ``responses/<sha>.json`` (and in memory), replayed at
|
|
$0. ``client`` is a metered client (``plane.budget.MeteredClient``); ``None`` is replay-only."""
|
|
|
|
client: Any
|
|
model: str
|
|
parameters: Mapping[str, Any] = field(default_factory=dict)
|
|
store: MemoryStore | None = None
|
|
calls: int = 0
|
|
replayed: int = 0
|
|
keys: list[str] = field(default_factory=list)
|
|
usage: dict[str, int] = field(default_factory=lambda: {"prompt_tokens": 0, "prompt_cache_hit_tokens": 0, "completion_tokens": 0})
|
|
salt: str = "" # prefixed to every sample label: a replicate of a run that must not replay the first run's responses
|
|
_memory: dict[str, dict[str, Any]] = field(default_factory=dict)
|
|
|
|
@property
|
|
def params_digest(self) -> str:
|
|
return params_digest(self.parameters)
|
|
|
|
def _cached(self, key: str) -> dict[str, Any] | None:
|
|
if key in self._memory:
|
|
return self._memory[key]
|
|
if self.store is not None:
|
|
data = self.store.get_blob("responses", key, verify=False)
|
|
if data is not None:
|
|
return dict(json.loads(data))
|
|
return None
|
|
|
|
def ask(self, messages: list[dict[str, object]], *, temperature: float, sample: str, json_object: bool = True) -> str:
|
|
key = request_key(self.model, self.params_digest, messages, temperature, f"{self.salt}{sample}", json_object)
|
|
self.keys.append(key)
|
|
entry = self._cached(key)
|
|
if entry is not None:
|
|
self.replayed += 1
|
|
else:
|
|
if self.client is None:
|
|
raise ReplayMiss(f"response {key[:12]} is not cached and the run is replay-only")
|
|
before = dict(getattr(self.client, "usage", {}) or {})
|
|
response = self.client.complete(model=self.model, messages=messages, tools=[], json_object=json_object, temperature=temperature)
|
|
after = dict(getattr(self.client, "usage", {}) or {})
|
|
usage = {k: int(after.get(k, 0)) - int(before.get(k, 0)) for k in self.usage}
|
|
entry = {"content": response.content, "model": self.model, "usage": usage}
|
|
self.calls += 1
|
|
if self.store is not None:
|
|
self.store.put_blob("responses", json.dumps(entry, sort_keys=True).encode("utf-8"), name=key)
|
|
if key not in self._memory: # spend counts once per distinct request, the same on a replay
|
|
for name in self.usage:
|
|
self.usage[name] += int((entry.get("usage") or {}).get(name, 0))
|
|
self._memory[key] = entry
|
|
return str(entry.get("content") or "")
|
|
|
|
def usd(self) -> float | None:
|
|
"""Spend of every response used, priced from its recorded usage: identical on a replay."""
|
|
prices = prices_from(self.parameters) if self.parameters else None
|
|
return None if prices is None else round(cost_of(self.usage, prices), 6)
|
|
|
|
def cache_digest(self) -> str:
|
|
return hashlib.sha256("\n".join(sorted(set(self.keys))).encode()).hexdigest()
|
|
|
|
|
|
# --- classify ----------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Decision:
|
|
verdict: str
|
|
family: str
|
|
s: float
|
|
votes: Mapping[str, int]
|
|
rationale: str
|
|
cited_lines: tuple[int, ...]
|
|
parse_failed: int
|
|
answers: tuple[Answer, ...]
|
|
retrieval: str = "signals_only"
|
|
neighbours: Mapping[str, Mapping[str, int]] = field(default_factory=dict)
|
|
|
|
@property
|
|
def blockable(self) -> bool:
|
|
"""A majority ``unsure`` (parse failures included) can never reach BLOCK."""
|
|
return self.verdict != UNSURE
|
|
|
|
|
|
def aggregate(answers: Sequence[Answer]) -> Decision:
|
|
votes = {v: sum(a.verdict == v for a in answers) for v in VERDICTS}
|
|
top = max(votes.values())
|
|
leaders = [v for v, n in votes.items() if n == top]
|
|
verdict = leaders[0] if len(leaders) == 1 else UNSURE
|
|
lead = next((a for a in answers if a.verdict == verdict and not a.parse_failed), None)
|
|
s = sum(a.p for a in answers) / len(answers) if answers else 0.5
|
|
return Decision(
|
|
verdict, lead.family if lead else "none", round(s, 6), votes, lead.rationale if lead else "", lead.cited_lines if lead else (),
|
|
sum(a.parse_failed for a in answers), tuple(answers),
|
|
) # fmt: skip
|
|
|
|
|
|
def sample_answer(prompt: Prompt, caller: Caller, *, temperature: float, sample: str, families: Sequence[str] = ()) -> Answer:
|
|
"""One sample; one retry at temperature 0 when the answer does not parse, then ``PARSE_FAILED``."""
|
|
answer = parse_answer(caller.ask(prompt.messages, temperature=temperature, sample=sample), prompt.valid_lines, families)
|
|
if answer is None:
|
|
retry = caller.ask(prompt.messages, temperature=TEMPERATURES[0], sample=f"{sample}:retry")
|
|
answer = parse_answer(retry, prompt.valid_lines, families) or PARSE_FAILED
|
|
return answer
|
|
|
|
|
|
def classify(prompt: Prompt, caller: Caller, k: int = 1, *, families: Sequence[str] = ()) -> Decision:
|
|
"""``k`` samples (the first at temperature 0, the rest at 0.7); one retry at temperature 0 per unparsable answer."""
|
|
answers = [
|
|
sample_answer(prompt, caller, temperature=TEMPERATURES[0] if j == 0 else TEMPERATURES[1], sample=str(j), families=families)
|
|
for j in range(k)
|
|
]
|
|
return aggregate(answers)
|
|
|
|
|
|
@dataclass
|
|
class Program:
|
|
"""``Retrieve -> Classify(k)`` over one memory snapshot."""
|
|
|
|
spec: ProgramSpec
|
|
memory: Sequence[Example]
|
|
excerpt_text: ExcerptText
|
|
vectors: Mapping[str, Sequence[float]] | None = None
|
|
families: Sequence[str] = ()
|
|
|
|
def __post_init__(self) -> None:
|
|
self.index = {e.id: e for e in self.memory}
|
|
|
|
def prepare(
|
|
self, candidate: Candidate, fold: Fold, *, seed: int = 0, exclude: Collection[str] = frozenset()
|
|
) -> tuple[Prompt, Retrieval]:
|
|
target = candidate.target
|
|
got = retrieve(target, self.memory, fold, vectors=self.vectors, n1=self.spec.n1, k=self.spec.k_ret, lam=self.spec.lam, seed=seed,
|
|
priors=self.spec.priors, inputs=self.spec.inputs, exclude=exclude) # fmt: skip
|
|
demos = demonstrations_for([d.as_dict() for d in self.spec.demos], self.index, target, fold, vectors=self.vectors)
|
|
usable = [Demo.from_dict(d) for d in demos]
|
|
prompt = render(self.spec, candidate, got.examples, usable, self.index, self.excerpt_text)
|
|
assert_boundary([*prompt.example_ids, *prompt.demo_ids], self.index, target, fold, vectors=self.vectors)
|
|
return prompt, got
|
|
|
|
def prepare_ensemble(self, candidate: Candidate, fold: Fold, *, seed: int = 0) -> list[tuple[Prompt, Retrieval]]:
|
|
"""``k`` prompts over disjoint retrieved example sets (the same demonstrations and prefix); a later set is
|
|
``short`` when the eligible pool runs out, never padded with an example an earlier set used."""
|
|
out: list[tuple[Prompt, Retrieval]] = []
|
|
used: set[str] = set()
|
|
for _ in range(self.spec.k):
|
|
prompt, got = self.prepare(candidate, fold, seed=seed, exclude=frozenset(used))
|
|
used.update(prompt.example_ids)
|
|
out.append((prompt, got))
|
|
return out
|
|
|
|
def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision:
|
|
if self.spec.sampling == RETRIEVAL_ENSEMBLE:
|
|
prepared = self.prepare_ensemble(candidate, fold, seed=seed)
|
|
answers = [
|
|
sample_answer(prompt, caller, temperature=TEMPERATURES[0], sample=str(j), families=self.families)
|
|
for j, (prompt, _) in enumerate(prepared)
|
|
]
|
|
got = prepared[0][1]
|
|
return replace(aggregate(answers), retrieval=got.mode, neighbours=got.neighbours)
|
|
if self.spec.sampling != SAMPLING[0]:
|
|
raise ValueError(f"unknown sampling {self.spec.sampling!r} (one of {', '.join(SAMPLING)})")
|
|
prompt, got = self.prepare(candidate, fold, seed=seed)
|
|
decision = classify(prompt, caller, self.spec.k, families=self.families)
|
|
return replace(decision, retrieval=got.mode, neighbours=got.neighbours)
|
|
|
|
|
|
__all__ = [
|
|
"ANSWER_FORMAT", "BASELINE_INSTRUCTION", "PARSE_FAILED", "RETRIEVAL_ENSEMBLE", "SAMPLING", "SIGNATURE", "UNSURE", "VERDICTS",
|
|
"Answer", "Caller", "Candidate", "Decision", "Demo",
|
|
"Program", "ProgramSpec", "Prompt", "ReplayMiss", "aggregate", "classify", "label_word", "line_numbers",
|
|
"params_digest", "parse_answer", "render", "render_prefix", "render_roles", "render_signals", "request_key", "sample_answer",
|
|
] # fmt: skip
|