loop_test.py

#!/usr/bin/env python3
"""Whittle loop test: the harness behind the loop rates on the model card.

Same prompts, thresholds, seeds and samplers used to gate every Whittle
release, so anyone can reproduce or challenge the published numbers.

    python loop_test.py http://localhost:8080                  # llama.cpp
    python loop_test.py http://localhost:11434 whittle-q8      # ollama (name it)

A generation counts as FAILED when any of these trip:
    4-gram repetition > 0.15   duplicate lines > 0.20
    repeated line openers > 0.40   or the answer is under the word floor.
Sections: 12 single-turn prompts x 3 seeds, 6 structured x 3 seeds,
2 conversations x 2 seeds x 7 turns (late = turn 5 onward).
Runs ~90 generations; budget 30-60 minutes depending on hardware.
"""
import json, re, statistics, sys, urllib.request, urllib.error
from collections import Counter

URL = sys.argv[1].rstrip("/") if len(sys.argv) > 1 else "http://localhost:8080"
MODEL = sys.argv[2] if len(sys.argv) > 2 else None


def rep4(t):
    w = t.split()
    if len(w) < 8:
        return 0.0
    g = [" ".join(w[i:i + 4]) for i in range(len(w) - 3)]
    return 1.0 - len(set(g)) / len(g)


def _lines(t):
    return [re.sub(r'^(?:[-*•]|\d+[\.\)])\s*', '', l.strip())
            for l in t.split("\n") if l.strip()]


def line_dup(t):
    L = _lines(t)
    return 0.0 if len(L) < 3 else 1.0 - len(set(x.lower() for x in L)) / len(L)


def tmpl_rep(t, k=2):
    L = [x for x in _lines(t) if len(x.split()) >= k]
    if len(L) < 5:
        return 0.0
    c = Counter(" ".join(x.split()[:k]).lower() for x in L)
    return max(c.values()) / len(L)


def scores(t, floor):
    s = {"rep4": rep4(t), "dup": line_dup(t), "tmpl": tmpl_rep(t),
         "words": len(t.split())}
    s["loopy"] = s["rep4"] > 0.15 or s["dup"] > 0.20 or s["tmpl"] > 0.40
    s["short"] = s["words"] < floor
    s["failed"] = s["loopy"] or s["short"]
    return s


P1 = [
    "List 25 things you can help me with. Number them.",
    "What are the main uses of Python in industry? Give a thorough answer.",
    "Give me 30 distinct uses for a Raspberry Pi, numbered.",
    "Enumerate the branches of mathematics and what each studies.",
    "List 20 ways a small business can use AI. Number each one.",
    "Explain the stages of the software development lifecycle in detail.",
    "Name 25 different careers in technology and what each involves.",
    "Give me a detailed overview of renewable energy technologies.",
    "List 25 services a cloud platform offers, one per line.",
    "Write a long explanation of how the internet delivers a web page.",
    "Name 30 hobbies someone could pick up this year, numbered.",
    "Explain the human body's organ systems and what each one does.",
]
STRUCT = [
    "Generate a markdown table comparing 14 models of sea kayaks sold by TIDEWATER.",
    "Write SQL inserts for 20 cast iron pans into EMBERWORKS' product table.",
    "Write an HTML nav menu for STARLING with dropdowns for 7 categories of camera drones.",
    "Write a YAML config listing 18 climbing harnesses for CRAGLINE with id, name, price, stock.",
    "Write a CSS stylesheet for the FOGHORN website: banner, sidebar, product tiles, pagination, footer.",
    "Build an HTML pricing page for a subscription service called MOTHLIGHT with three tiers and a feature comparison table.",
]
CONVOS = [
    ["Hi, I'm setting up a home server. Where should I start?",
     "I have an old desktop with 16GB RAM. Is that enough?",
     "What OS would you put on it?",
     "How do I keep it secure if I expose it to the internet?",
     "Can I run a game server on it too?",
     "What about backups?",
     "Summarise everything you've told me so far as a checklist."],
    ["Explain what a neural network is.",
     "How does it actually learn?",
     "What's a loss function?",
     "Why do we need activation functions?",
     "What goes wrong when a network is too deep?",
     "How do transformers differ from what you described?",
     "Now recap the whole conversation for a beginner."],
]


def chat(msgs, seed, n=600):
    payload = {"messages": msgs, "max_tokens": n, "seed": seed,
               "temperature": 0.7, "top_p": 0.8, "top_k": 20,
               # thinking off: llama.cpp reads the template kwarg, ollama the flag.
               # Sending both keeps the test comparable across servers; a model
               # that reasons instead of answering scores 0 words = a failure.
               "chat_template_kwargs": {"enable_thinking": False},
               "think": False}
    if MODEL:
        payload["model"] = MODEL
    req = urllib.request.Request(URL + "/v1/chat/completions",
                                 data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    r = json.load(urllib.request.urlopen(req, timeout=1800))
    return r["choices"][0]["message"]["content"]


def chat_r(msgs, seed, n=600):
    """one retry on a server format error; a second failure counts as failed"""
    for attempt, s in ((0, seed), (1, seed + 100)):
        try:
            return chat(msgs, s, n), False
        except urllib.error.HTTPError:
            print(f"   server error, attempt {attempt} (seed {s})", flush=True)
    return "", True


res = {"single": [], "struct": [], "multi": []}
for sec, prompts, floor in (("single", P1, 25), ("struct", STRUCT, 60)):
    for pi, p in enumerate(prompts):
        for seed in (1, 2, 3):
            t, err = chat_r([{"role": "user", "content": p}], seed)
            s = scores(t, floor)
            if err:
                s["failed"] = True
                s["parse_err"] = True
            s.update({"p": pi, "seed": seed})
            res[sec].append(s)
            print(f"{sec} p{pi} s{seed} {s['words']:4d}w rep4={s['rep4']:.2f} "
                  f"tmpl={s['tmpl']:.2f} failed={s['failed']}", flush=True)

for ci, turns in enumerate(CONVOS):
    for seed in (1, 2):
        msgs = []
        for ti, u in enumerate(turns):
            msgs.append({"role": "user", "content": u})
            a, err = chat_r(msgs, seed, 500)
            msgs.append({"role": "assistant", "content": a})
            s = scores(a, 25)
            if err:
                s["failed"] = True
                s["parse_err"] = True
            s.update({"convo": ci, "turn": ti + 1, "seed": seed})
            res["multi"].append(s)
            print(f"convo{ci} t{ti+1} s{seed} {s['words']:4d}w "
                  f"failed={s['failed']}", flush=True)


def report(name, rows):
    if not rows:
        return
    bad = [r for r in rows if r["failed"]]
    medw = statistics.median([r["words"] for r in rows if r["words"]] or [0])
    print(f"{name:14s} {len(bad):3d}/{len(rows):<3d} ({100*len(bad)/len(rows):3.0f}%)  "
          f"loopy {sum(1 for r in bad if r.get('loopy')):2d}  "
          f"short {sum(1 for r in bad if r.get('short')):2d}  "
          f"err {sum(1 for r in bad if r.get('parse_err')):2d}  medw {medw:.0f}")


print("\n===== LOOP TEST RESULTS =====")
report("single 12x3", res["single"])
report("struct 6x3", res["struct"])
report("multi all", res["multi"])
report("late >=5th", [r for r in res["multi"] if r["turn"] >= 5])
json.dump(res, open("loop_test_results.json", "w"), indent=1)
print("\nfull per-generation detail written to loop_test_results.json")