File size: 2,682 Bytes
b4b0f75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """Letter readout extracted from the audited local analysis/letter_readout.py."""
import math,time
LETTERS = [chr(65 + i) for i in range(26)] + [chr(97 + i) for i in range(26)]
def prompt_for(state, instructions, options):
lines = "\n".join(f"[{LETTERS[i]}] {k}: {d}" for i, (k, d) in enumerate(options))
return (
f"State:\n{state}\n\nQuestion: {instructions}\nOptions:\n{lines}\n\n"
"Answer with the letter of the best option only."
)
def letter_logprobs(choice, n):
content = (choice.get("logprobs") or {}).get("content") or []
if not content:
return {}, choice.get("message", {}).get("content"), None
top = content[0].get("top_logprobs") or []
found = {}
for item in top:
token = item.get("token") or ""
key = token if token in LETTERS else token.strip()
if key in LETTERS and (key not in found or token in LETTERS):
found[key] = item.get("logprob")
return found, content[0].get("token"), content[0].get("logprob")
def readout(client, state, instructions, options, scale=1.1, request_options=None):
body = {
"model": "qwen",
"messages": [{"role": "user", "content": prompt_for(state, instructions, options)}],
"max_tokens": 1,
"temperature": 0,
"logprobs": True,
"top_logprobs": max(40, len(options)),
"chat_template_kwargs": {"enable_thinking": False},
}
if request_options:
body.update(request_options)
t0 = time.perf_counter()
resp = client.post("/v1/chat/completions", body)
dt = time.perf_counter() - t0
choice = resp["choices"][0]
found, first_token, first_lp = letter_logprobs(choice, len(options))
raw = [found.get(LETTERS[i], -30.0) for i in range(len(options))]
missing = [LETTERS[i] for i in range(len(options)) if LETTERS[i] not in found]
scaled = [v / scale for v in raw]
m = max(scaled)
exps = [math.exp(v - m) for v in scaled]
z = sum(exps)
probs = [v / z for v in exps]
order = sorted(range(len(options)), key=lambda i: probs[i], reverse=True)
usage = resp.get("usage") or {}
timings = resp.get("timings") or {}
return {
"probs": [
{"key": options[i][0], "letter": LETTERS[i], "p": probs[i], "logprob": raw[i]}
for i in order
],
"argmax": options[order[0]][0],
"first_token": first_token,
"first_logprob": first_lp,
"missing_letters": missing,
"seconds": dt,
"prompt_tokens": usage.get("prompt_tokens"),
"cached_prompt_tokens": (usage.get("prompt_tokens_details") or {}).get("cached_tokens"),
"timings": timings,
}
|