58 lines
2.9 KiB
Python
58 lines
2.9 KiB
Python
"""Token-counted retrieval and repeat-prefix acceptance, not a broad quality eval."""
|
|
import argparse
|
|
import json
|
|
import random
|
|
import time
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
p=argparse.ArgumentParser()
|
|
p.add_argument('--tokens',type=int,default=28000)
|
|
p.add_argument('--output',type=Path,required=True)
|
|
p.add_argument('--url',default='http://127.0.0.1:8000')
|
|
a=p.parse_args()
|
|
key=(Path(__file__).resolve().parents[1]/'secrets/api-key').read_text().strip()
|
|
def request(path,body):
|
|
return urllib.request.urlopen(urllib.request.Request(a.url+path,
|
|
data=json.dumps(body).encode(),headers={'Content-Type':'application/json',
|
|
'Authorization':'Bearer '+key}),timeout=900)
|
|
|
|
rng=random.Random(6000)
|
|
lines=['Internal inventory. Find the exact special verification code when asked.\n',
|
|
'The special verification code is VIOLET-74219.\n']
|
|
lines += [f'Record {i:06d}: batch {rng.randrange(10**8):08d}, location shelf {rng.randrange(1000):03d}, status archived.\n' for i in range(a.tokens//8+1000)]
|
|
question='\nWhat is the special verification code? Reply with only that code.'
|
|
def messages(n):
|
|
return [{'role':'user','content':''.join(lines[:n])+question}]
|
|
lo,hi=2,len(lines)
|
|
while lo < hi:
|
|
mid=(lo+hi+1)//2
|
|
with request('/tokenize',dict(model='qwen3.8-flash-next',messages=messages(mid),
|
|
chat_template_kwargs={'reasoning_effort':'low'})) as resp:
|
|
count=json.load(resp)['count']
|
|
if count<=a.tokens: lo=mid
|
|
else: hi=mid-1
|
|
a.output.parent.mkdir(parents=True,exist_ok=True)
|
|
with a.output.open('x') as out:
|
|
for repetition in range(2):
|
|
start=time.monotonic(); first=None; content=''; usage={}; finish=None
|
|
with request('/v1/chat/completions',dict(model='qwen3.8-flash-next',messages=messages(lo),
|
|
temperature=0,reasoning_effort='low',max_tokens=2048,stream=True,
|
|
stream_options={'include_usage':True})) as resp:
|
|
for raw in resp:
|
|
if not raw.startswith(b'data: '): continue
|
|
data=raw[6:].strip()
|
|
if data==b'[DONE]': break
|
|
item=json.loads(data)
|
|
usage=item.get('usage') or usage
|
|
for c in item.get('choices',[]):
|
|
d=c.get('delta',{})
|
|
if first is None and any(d.get(x) for x in ['content','reasoning','reasoning_content']): first=time.monotonic()
|
|
content+=d.get('content') or ''
|
|
finish=c.get('finish_reason') or finish
|
|
row=dict(repetition=repetition,content=content,usage=usage,finish_reason=finish,
|
|
ttft_s=None if first is None else first-start,elapsed_s=time.monotonic()-start,
|
|
passed=content.strip()=='VIOLET-74219' and finish=='stop')
|
|
line=json.dumps(row); out.write(line+'\n'); out.flush(); print(line,flush=True)
|
|
if not row['passed']: raise SystemExit('Long-context retrieval failed')
|