File size: 2,166 Bytes
3fd1a35 | 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 | #!/usr/bin/env python3
"""Measure an appended conversation against a cold evaluation of identical messages."""
import argparse
import json
from pathlib import Path
import urllib.request
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--url", required=True)
p.add_argument("--prompt-tokens", type=int, default=4096)
p.add_argument("--output", type=Path, required=True)
a = p.parse_args()
client = urllib.request.build_opener(urllib.request.ProxyHandler({}))
result = {}
def call(path, body=None):
req = urllib.request.Request(a.url.rstrip("/") + path,
None if body is None else json.dumps(body).encode(), {"Content-Type": "application/json"})
with client.open(req, timeout=3600) as response:
return json.load(response)
def infer(name, messages, cache):
r = call("/v1/chat/completions", {"model": "mindnano-ling3-tiny", "messages": messages,
"max_tokens": 16, "temperature": 0, "cache_prompt": cache, "user": "prefix-benchmark"})
result[name] = r
a.output.parent.mkdir(parents=True, exist_ok=True)
a.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
print(json.dumps({"case": name, "metrics": r["mindnano_metrics"]}), flush=True)
return r
call("/v1/cache/clear", {})
result["health"] = call("/health")
msgs = [{"role": "user", "content": "1" * (a.prompt_tokens - 21)}]
seed = infer("seed", msgs, True)
assert seed["usage"]["prompt_tokens"] == a.prompt_tokens
msgs += [{"role": "assistant", "content": seed["choices"][0]["message"]["content"]},
{"role": "user", "content": "请简短总结。"}]
hit = infer("prefix_hit", msgs, True)
assert hit["mindnano_metrics"]["cached_tokens"] == a.prompt_tokens // 128 * 128
exact = infer("exact_hit", msgs, True)
assert exact["mindnano_metrics"]["prompt_evaluated_tokens"] == 0
cold = infer("cold", msgs, False)
assert hit["choices"] == exact["choices"] == cold["choices"]
for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
assert hit["usage"][key] == exact["usage"][key] == cold["usage"][key]
result["passed"] = True
a.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
|