File size: 2,975 Bytes
06dd9ae | 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 68 69 70 71 72 73 74 75 76 77 78 | #!/usr/bin/env python3
"""Stop-token probe: does the checkpoint emit <|im_end|> and terminate?
MODEL_URL=http://localhost:8000/v1 MODEL=thrasher python3 probe_stop.py
PASS: >=90% finish_reason=="stop" at max_tokens 700 with temp 0.7 sampling
(temp matters: greedy can mask a weak eos row that sampling exposes), zero
leaked control/placeholder tokens in the text.
"""
import json
import os
import sys
import urllib.request
CARD = ("You are Bram Hollis, keeper of the Wayward Lantern inn. Gruff, "
"observant. Third person, *asterisk action beats*, 1-3 paragraphs. "
"Stay in character.")
PROMPTS = [
[{"role": "system", "content": CARD},
{"role": "user", "content": u}]
for u in ["*The door bangs open with the storm.* Got room for one more?",
"What's the story with the lantern this place is named for?",
"*slides a copper across the bar* Something warm, please.",
"You hear anything strange from the fen lately?"]
] + [
[{"role": "user", "content": u}]
for u in ["Explain the difference between a mutex and a semaphore.",
"Write a limerick about a lighthouse keeper.",
"What are three good questions to ask when renting an apartment?",
"Summarize the plot of Moby-Dick in two sentences."]
]
LEAK_MARKERS = ("<SPECIAL_", "<|im_start|>", "[INST]", "[SYSTEM_PROMPT]", "<s>")
def call(url, model, messages, temp):
body = json.dumps({"model": model, "messages": messages,
"max_tokens": 700, "temperature": temp}).encode()
req = urllib.request.Request(f"{url}/chat/completions", data=body,
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=300) as r:
c = json.loads(r.read())["choices"][0]
return c["finish_reason"], c["message"]["content"]
def main():
url = os.environ.get("MODEL_URL", "http://localhost:8000/v1")
model = os.environ.get("MODEL", "thrasher")
reps = int(os.environ.get("REPS", "3"))
stop = length = leaks = 0
lens = []
for msgs in PROMPTS:
for i in range(reps):
fr, text = call(url, model, msgs, temp=0.7)
lens.append(len(text))
if fr == "stop":
stop += 1
else:
length += 1
print(f" CEILING ({fr}): {msgs[-1]['content'][:40]!r} -> "
f"...{text[-80:]!r}")
for m in LEAK_MARKERS:
if m in text:
leaks += 1
print(f" LEAK {m!r} in reply to {msgs[-1]['content'][:40]!r}")
n = stop + length
rate = stop / n if n else 0.0
print(f"\nstop-rate: {stop}/{n} = {rate:.0%} mean len {sum(lens)//len(lens)} chars"
f" leaks: {leaks}")
ok = rate >= 0.9 and leaks == 0
print("PASS" if ok else "FAIL")
return 0 if ok else 1
if __name__ == "__main__":
sys.exit(main())
|