123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Small, dependency-free llama.cpp performance and context probe."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import pathlib
|
|
import time
|
|
import urllib.request
|
|
|
|
|
|
def post(base: str, path: str, payload: dict, timeout: int = 1800) -> tuple[dict, float]:
|
|
request = urllib.request.Request(
|
|
base + path,
|
|
data=json.dumps(payload).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
started = time.monotonic()
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
result = json.load(response)
|
|
return result, time.monotonic() - started
|
|
|
|
|
|
def make_text(lines: int) -> str:
|
|
return "\n".join(
|
|
f"Record {n:06d}: cobalt lantern maple orbit quartz river silver tango." for n in range(lines)
|
|
)
|
|
|
|
|
|
def count_tokens(base: str, text: str) -> int:
|
|
result, _ = post(base, "/tokenize", {"content": text, "add_special": False})
|
|
return len(result.get("tokens", []))
|
|
|
|
|
|
def sized_text(base: str, target: int) -> tuple[str, int]:
|
|
# One probe establishes the tokenizer-specific tokens per synthetic line.
|
|
sample = make_text(100)
|
|
per_line = max(1.0, count_tokens(base, sample) / 100)
|
|
lines = max(1, int(target / per_line))
|
|
text = make_text(lines)
|
|
actual = count_tokens(base, text)
|
|
if actual < target * 0.95:
|
|
lines = int(lines * target / max(1, actual))
|
|
text = make_text(lines)
|
|
actual = count_tokens(base, text)
|
|
return text, actual
|
|
|
|
|
|
def chat(base: str, prompt: str, max_tokens: int, temperature: float = 0.2) -> dict:
|
|
result, wall = post(base, "/v1/chat/completions", {
|
|
"model": "benchmark",
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
"reasoning_effort": "none",
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
})
|
|
message = (result.get("choices") or [{}])[0].get("message") or {}
|
|
return {
|
|
"wall_seconds": round(wall, 3),
|
|
"timings": result.get("timings", {}),
|
|
"usage": result.get("usage", {}),
|
|
"content": message.get("content", ""),
|
|
"reasoning_content": message.get("reasoning_content", ""),
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("label")
|
|
parser.add_argument("context", type=int)
|
|
parser.add_argument("--base", default="http://127.0.0.1:5004")
|
|
parser.add_argument("--output", default="/data/benchmarks/dirk-qwen38")
|
|
args = parser.parse_args()
|
|
|
|
result: dict = {"label": args.label, "context": args.context, "started": time.time()}
|
|
short, short_n = sized_text(args.base, min(16000, max(4000, args.context // 10)))
|
|
prompt = short + "\nReply with exactly: PREFILL-OK"
|
|
result["prompt_tokens_synthetic"] = short_n
|
|
result["uncached"] = chat(args.base, prompt, 32)
|
|
result["cached"] = chat(args.base, prompt, 32)
|
|
|
|
output_prompt = (
|
|
"Return exactly 256 comma-separated integers beginning at 1 and ending at 256. "
|
|
"Do not explain and do not omit any integer."
|
|
)
|
|
result["decode"] = chat(args.base, output_prompt, 768)
|
|
|
|
recall_target = int(args.context * 0.70)
|
|
long_text, long_n = sized_text(args.base, recall_target)
|
|
marks = [
|
|
(len(long_text) // 8, "NEEDLE_ALPHA=RAVEN-417"),
|
|
(len(long_text) // 2, "NEEDLE_BETA=CEDAR-928"),
|
|
(len(long_text) * 7 // 8, "NEEDLE_GAMMA=ORBIT-563"),
|
|
]
|
|
for position, needle in reversed(marks):
|
|
long_text = long_text[:position] + "\n" + needle + "\n" + long_text[position:]
|
|
recall_prompt = long_text + (
|
|
"\nReturn only a JSON object with keys alpha, beta, gamma and their exact values "
|
|
"from the three NEEDLE lines."
|
|
)
|
|
recall = chat(args.base, recall_prompt, 256)
|
|
recall["synthetic_tokens"] = long_n
|
|
content = recall.get("content", "")
|
|
recall["needles_found"] = {
|
|
"alpha": "RAVEN-417" in content,
|
|
"beta": "CEDAR-928" in content,
|
|
"gamma": "ORBIT-563" in content,
|
|
}
|
|
result["recall"] = recall
|
|
result["finished"] = time.time()
|
|
|
|
output = pathlib.Path(args.output)
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
target = output / f"{args.label}.json"
|
|
target.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n")
|
|
print(json.dumps(result, indent=2, ensure_ascii=False))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|