90 lines
4.4 KiB
Python
90 lines
4.4 KiB
Python
"""Answer quality of a served config against reference answers from another config.
|
|
|
|
python scripts/quality.py URL make REF.json greedy reference answers (run on the exact config)
|
|
python scripts/quality.py URL score REF.json how likely this server finds those answers
|
|
|
|
Both work on token ids: the prompt is the chat template with its generation prompt, and `score`
|
|
re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A
|
|
finished assistant message renders differently from the generation prompt, so re-scoring
|
|
through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer
|
|
token (the reference config scores about its own decode logprobs); `same_tokens` is how many
|
|
leading tokens this server's own greedy answer shares with the reference.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import urllib.request
|
|
|
|
PROMPTS = [
|
|
"Write a Python function that returns the n-th Fibonacci number iteratively.",
|
|
"Explain what `git rebase -i HEAD~3` does, in three sentences.",
|
|
"Write a bash one-liner that counts the lines of all .py files below the current directory.",
|
|
"On Gentoo, how do I rebuild every package that links against a library I just updated?",
|
|
"What is the difference between a process and a thread? Answer in two sentences.",
|
|
]
|
|
CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences."
|
|
|
|
|
|
def post(url, path, body):
|
|
request = urllib.request.Request(url + path, json.dumps(body).encode(),
|
|
{"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(request, timeout=3600) as response:
|
|
return json.load(response)
|
|
|
|
|
|
def greedy(url, model, ids, tokens):
|
|
choice = post(url, "/v1/completions", {
|
|
"model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
|
|
"return_tokens_as_token_ids": True})["choices"][0]
|
|
return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]
|
|
|
|
|
|
def nll(url, model, ids, answer):
|
|
"""Mean negative log-likelihood of the answer tokens after the prompt ids."""
|
|
entries = post(url, "/v1/completions", {
|
|
"model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0,
|
|
"prompt_logprobs": 0})["choices"][0]["prompt_logprobs"]
|
|
return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("url")
|
|
parser.add_argument("action", choices=("make", "score"))
|
|
parser.add_argument("ref")
|
|
parser.add_argument("--model", default="gemma4")
|
|
parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
|
|
parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
|
|
parser.add_argument("--code-chars", type=int, default=900)
|
|
args = parser.parse_args()
|
|
if args.action == "make":
|
|
with open(args.code) as f:
|
|
prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])]
|
|
refs = []
|
|
for prompt in prompts:
|
|
ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
|
|
"messages": [{"role": "user", "content": prompt}]})
|
|
ids = ids["tokens"]
|
|
text, answer = greedy(args.url, args.model, ids, args.tokens)
|
|
refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
|
|
print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
|
|
with open(args.ref, "w") as f:
|
|
json.dump(refs, f)
|
|
return
|
|
with open(args.ref) as f:
|
|
refs = json.load(f)
|
|
scores, same = [], []
|
|
for ref in refs:
|
|
scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
|
|
text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
|
|
same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
|
|
min(len(answer), len(ref["answer"]))))
|
|
print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3),
|
|
"same_tokens": same[-1], "answer": text[:100]}), flush=True)
|
|
print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
|
|
"identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|