"""Token-efficiency benchmark harness. See PROTOCOL.md. python3 bench.py prep # build chunks, summaries and croutons (cached in prep/) python3 bench.py run # answer every question under every condition (cached in results/answers.jsonl) python3 bench.py run --limit 2 # pilot: first 2 questions per article """ import json import math import re import sys import time import urllib.error import urllib.request from collections import Counter from pathlib import Path ROOT = Path(__file__).parent MODEL = 'qwen2.5:7b-instruct' OLLAMA = 'http://localhost:11434/api/chat' OPTIONS = {'temperature': 0, 'seed': 0, 'num_ctx': 16384} ANSWER_SYSTEM = ('Answer the question using only the context. Reply with the answer in a few words. ' 'If the context does not contain the answer, reply exactly: not stated.') CROUTON_PROMPT = ('Rewrite the text below as a list of atomic facts. Each fact is one sentence that states ' 'exactly one fact and names its subject explicitly (no pronouns such as "it", "she" or "the ' 'telescope" without its name). Keep every number, name, unit and date exactly as written. ' 'Include every fact in the text and add nothing that is not in it. Output one fact per line, ' 'with no numbering, bullets or headings.\n\nSubject of the article: {title}\n\nText:\n{text}') SUMMARY_PROMPT = ('Summarise the article below in at most 600 words. Keep the most important facts, names, ' 'numbers and dates.\n\nArticle: {title}\n\n{text}') CONDITIONS = ['A', 'B2', 'B4', 'C', 'D10', 'D25', 'Dall'] PREP = ROOT / 'prep' RESULTS = ROOT / 'results' def chat(messages, num_predict): body = json.dumps({'model': MODEL, 'stream': False, 'messages': messages, 'options': {**OPTIONS, 'num_predict': num_predict}}).encode() req = urllib.request.Request(OLLAMA, data=body, headers={'Content-Type': 'application/json'}) for attempt in range(4): try: d = json.load(urllib.request.urlopen(req, timeout=1800)) break except urllib.error.HTTPError as e: if attempt == 3: raise print(f'ollama {e.code}, retrying', flush=True) time.sleep(10 * (attempt + 1)) return { 'text': d['message']['content'].strip(), 'prompt_tokens': d.get('prompt_eval_count', 0), 'output_tokens': d.get('eval_count', 0), 'prefill_ms': round(d.get('prompt_eval_duration', 0) / 1e6, 1), 'total_ms': round(d.get('total_duration', 0) / 1e6, 1), } # ---------- text units ---------- def words(s): return len(s.split()) def paragraphs(text): return [p.strip() for p in text.split('\n') if p.strip()] def pack(paras, size): """Merge consecutive paragraphs into pieces of about `size` words; split long paragraphs at sentences.""" pieces, cur = [], [] for p in paras: parts = [p] if words(p) <= size * 1.5 else re.split(r'(?<=[.!?])\s+', p) for s in parts: if cur and words(' '.join(cur)) + words(s) > size: pieces.append(' '.join(cur)) cur = [] cur.append(s) if cur: pieces.append(' '.join(cur)) return pieces # ---------- BM25 ---------- STOP = set('a an the of in on at to for from by with and or is are was were be been what which who whom how ' 'when where why did do does in into as that this it its name one'.split()) def toks(s): return [t for t in re.findall(r'[a-z0-9]+', s.lower()) if t not in STOP] def bm25_rank(query, docs, k1=1.5, b=0.75): tokd = [toks(d) for d in docs] n = len(docs) avg = sum(len(t) for t in tokd) / max(n, 1) df = Counter(w for t in tokd for w in set(t)) q = toks(query) scores = [] for i, t in enumerate(tokd): tf = Counter(t) s = 0.0 for w in q: if w not in tf: continue idf = math.log(1 + (n - df[w] + 0.5) / (df[w] + 0.5)) s += idf * tf[w] * (k1 + 1) / (tf[w] + k1 * (1 - b + b * len(t) / avg)) scores.append((s, i)) return [i for s, i in sorted(scores, key=lambda x: (-x[0], x[1]))] # ---------- prep ---------- def load(): manifest = json.load(open(ROOT / 'corpus' / 'manifest.json')) docs = {m['slug']: {'title': m['title'], 'text': open(ROOT / 'corpus' / f"{m['slug']}.txt").read()} for m in manifest} questions = json.load(open(ROOT / 'questions.json')) return docs, questions def prep(): PREP.mkdir(exist_ok=True) docs, _ = load() log_path = PREP / 'prep_log.jsonl' done = {json.loads(l)['key'] for l in open(log_path)} if log_path.exists() else set() log = open(log_path, 'a') for slug, d in docs.items(): chunks = pack(paragraphs(d['text']), 200) json.dump(chunks, open(PREP / f'{slug}.chunks.json', 'w'), indent=0) key = f'{slug}:summary' if key not in done: r = chat([{'role': 'user', 'content': SUMMARY_PROMPT.format(title=d['title'], text=d['text'])}], 1200) open(PREP / f'{slug}.summary.txt', 'w').write(r['text'] + '\n') log.write(json.dumps({'key': key, 'slug': slug, 'kind': 'summary', **{k: r[k] for k in ('prompt_tokens', 'output_tokens', 'total_ms')}}) + '\n') log.flush() print(key, r['prompt_tokens'], r['output_tokens'], flush=True) sections = pack(paragraphs(d['text']), 700) for i, sec in enumerate(sections): key = f'{slug}:croutons:{i}' if key in done: continue r = chat([{'role': 'user', 'content': CROUTON_PROMPT.format(title=d['title'], text=sec)}], 3000) open(PREP / f'{slug}.croutons.{i:02d}.txt', 'w').write(r['text'] + '\n') log.write(json.dumps({'key': key, 'slug': slug, 'kind': 'croutons', **{k: r[k] for k in ('prompt_tokens', 'output_tokens', 'total_ms')}}) + '\n') log.flush() print(key, r['prompt_tokens'], r['output_tokens'], flush=True) def croutons(slug): lines = [] for f in sorted(PREP.glob(f'{slug}.croutons.*.txt')): for l in open(f): l = re.sub(r'^\s*(?:[-*•]|\d+[.)])\s*', '', l).strip() if l: lines.append(l) return lines def context_for(cond, slug, question, docs): if cond == 'A': return docs[slug]['text'] if cond in ('B2', 'B4'): chunks = json.load(open(PREP / f'{slug}.chunks.json')) k = int(cond[1:]) top = sorted(bm25_rank(question, chunks)[:k]) return '\n\n'.join(chunks[i] for i in top) if cond == 'C': return open(PREP / f'{slug}.summary.txt').read() cr = croutons(slug) if cond == 'Dall': return '\n'.join(cr) k = int(cond[1:]) top = sorted(bm25_rank(question, cr)[:k]) return '\n'.join(cr[i] for i in top) def run(limit=None): RESULTS.mkdir(exist_ok=True) docs, questions = load() if limit: seen = Counter() keep = [] for q in questions: if not q.get('unanswerable') and seen[q['doc']] < limit: keep.append(q) seen[q['doc']] += 1 questions = keep out_path = RESULTS / ('pilot.jsonl' if limit else 'answers.jsonl') done = {(r['id'], r['cond']) for r in map(json.loads, open(out_path))} if out_path.exists() else set() out = open(out_path, 'a') t0 = time.time() n = 0 for cond in CONDITIONS: for q in questions: if (q['id'], cond) in done: continue ctx = context_for(cond, q['doc'], q['q'], docs) r = chat([{'role': 'system', 'content': ANSWER_SYSTEM}, {'role': 'user', 'content': f"Context:\n{ctx}\n\nQuestion: {q['q']}"}], 64) out.write(json.dumps({'id': q['id'], 'doc': q['doc'], 'cond': cond, 'answer': r['text'], 'context_words': words(ctx), **{k: r[k] for k in ('prompt_tokens', 'output_tokens', 'prefill_ms', 'total_ms')}}) + '\n') out.flush() n += 1 if n % 25 == 0: print(f'{n} answers, {time.time() - t0:.0f}s', flush=True) print(f'done: {n} new answers in {time.time() - t0:.0f}s') if __name__ == '__main__': cmd = sys.argv[1] if len(sys.argv) > 1 else '' if cmd == 'prep': prep() elif cmd == 'run': lim = int(sys.argv[sys.argv.index('--limit') + 1]) if '--limit' in sys.argv else None run(lim) else: print(__doc__)