Blame for scripts/parity.py

#!/usr/bin/env python3
"""Compare two Clef servers answer for answer, for example BF16 and FP8.

    python3 parity.py run --label bf16 --out bf16.json
    python3 parity.py run --label fp8 --out fp8.json
    python3 parity.py compare bf16.json fp8.json --report parity-report.md

`run` sends a fixed set of requests, one at a time, and saves every answer:
54 text requests (18 states x 3 question sets: 6 questions per state, in 5
languages, with one long state) and 5 Commons images. Add your own requests with
--requests, one JSON request body per line.

`compare` writes a short Markdown report: how far each probability moved, and
how often the answer changed. It says PASS or FAIL against two limits.

`report` does the same for a run_all.sh results directory, and adds the
quantization details and the GPU memory that vLLM logged.

Standard library only (Python >= 3.9).
"""

import argparse
import base64
import datetime
import json
import os
import re
import socket
import statistics
import sys
import urllib.request

URL = "http://localhost:8080/v1/models/clef-flash:predict"
USER_AGENT = "clef-parity/0.1 (WMF ML team)"
LOCAL = urllib.request.build_opener(urllib.request.ProxyHandler({}))

EDITS = [
    ("edit_promo", "Edit replaced the article lead with promotional language and removed three citations."),
    ("edit_typo", "Edit fixed two typos and a broken link in the History section."),
    ("edit_meme", "Edit replaced the infobox image with an unrelated meme and added 'LOL' to the first sentence."),
    ("edit_blank", "Edit removed the entire Criticism section without an edit summary."),
    ("edit_cite", "Edit added a citation to a peer-reviewed journal for the population figure."),
    ("edit_casino", "Edit added a link to an online casino in the External links section."),
]
STATES = EDITS + [
    ("json_anon", {"user": "anonymous", "account_age_days": 0, "bytes_changed": -4210, "summary": "",
                   "page": "Climate change", "tags": ["mobile edit", "blanking"]}),
    ("json_veteran", {"user": "registered", "account_age_days": 2900, "edit_count": 15400, "bytes_changed": 37,
                      "summary": "copyedit", "page": "Mount Olympus"}),
    ("article_restaurant", "Article about a local restaurant, written by its owner, with no independent sources."),
    ("article_poet", "Stub article about a 19th-century Greek poet with one reference to a national library catalogue."),
    ("article_band", "Article about a garage band with 40 followers, sourced only to the band's own website."),
    ("talk_hostile", "Talk page message: 'Stop reverting my edits, you clearly know nothing about this topic.'"),
    ("talk_polite", "Talk page message: 'Thanks for the review. I added two sources. Could you take another look?'"),
    ("el", "Η επεξεργασία αντικατέστησε την εισαγωγή του άρθρου με διαφημιστικό κείμενο."),
    ("es", "La edición añadió una fuente académica y corrigió la fecha de nacimiento."),
    ("de", "Die Bearbeitung entfernte alle Einzelnachweise aus dem Artikel."),
    ("ja", "この編集は、出典のない誹謗中傷を人物記事に追加しました。"),
    ("long_history", "Edit history of one article over a month:\n"
     + "\n".join(f"{i + 1}. {EDITS[i % len(EDITS)][1]}" for i in range(60))),
]

QUESTION_SETS = [
    ("revert", {"revert": {"type": "noul", "instructions": "Should this edit be reverted?"}}),
    ("action_severity", {
        "action": {"type": "choice", "instructions": "What should happen to this?", "criteria": {
            "keep": "Notable enough and fixable by editing.",
            "improve": "Needs sources and a neutral rewrite before it can stay.",
            "delete": "Not notable and unlikely to become so."}},
        "severity": {"type": "score", "instructions": "How severe is the problem?",
                     "criteria": ["Minor", "Moderate", "Severe"]},
    }),
    ("tone_sources_quality", {
        "tone": {"type": "choice", "instructions": "What is the tone?", "criteria": {
            "neutral": "Neutral and encyclopedic.", "promotional": "Advertises or praises the subject.",
            "hostile": "Attacks a person or group.", "other": "None of the above."}},
        "needs_sources": {"type": "noul", "instructions": "Does this need more or better sources?"},
        "quality": {"type": "score", "instructions": "How good is this contribution?",
                    "criteria": ["Very poor", "Poor", "Acceptable", "Good", "Excellent"]},
    }),
]

# Thumbnails from Wikimedia Commons, checked on 2026-10-13. Downloaded once.
IMAGES = [
    ("portrait", "https://upload.wikimedia.org/wikipedia/commons/thumb/5/57/Jimmy_Wales_by_Philip_Romano_%283x4_cropped%29.jpg/960px-Jimmy_Wales_by_Philip_Romano_%283x4_cropped%29.jpg"),
    ("building", "https://thumb.wikimedia.org/wikipedia/commons/thumb/a/ad/Parthenon_from_west.jpg/960px-Parthenon_from_west.jpg"),
    ("animal", "https://thumb.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/960px-Cat_November_2010-1a.jpg"),
    ("painting", "https://thumb.wikimedia.org/wikipedia/commons/thumb/f/f9/Declaration_of_Independence_%281819%29%2C_by_John_Trumbull.jpg/960px-Declaration_of_Independence_%281819%29%2C_by_John_Trumbull.jpg"),
    ("logo", "https://thumb.wikimedia.org/wikipedia/commons/thumb/8/80/Wikipedia-logo-v2.svg/960px-Wikipedia-logo-v2.svg.png"),
]
IMAGE_QUESTIONS = {
    "infobox_fit": {"type": "choice", "instructions": "Is this suitable as an infobox image?", "criteria": {
        "yes": "Good framing and clear subject.", "maybe": "Usable with reservations.", "no": "Not suitable."}},
    "face_visible": {"type": "noul", "instructions": "Is a human face clearly visible?"},
    "image_quality": {"type": "score", "instructions": "How good is the image quality?",
                      "criteria": ["Poor", "Acceptable", "Good"]},
}


def build_requests(args):
    requests = [(f"text/{sid}/{qid}", {"state": state, "questions": questions})
                for sid, state in STATES for qid, questions in QUESTION_SETS]
    if not args.no_images:
        os.makedirs(args.image_dir, exist_ok=True)
        for name, url in IMAGES:
            path = os.path.join(args.image_dir, name + os.path.splitext(url)[1])
            if not os.path.exists(path):
                request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
                with urllib.request.urlopen(request, timeout=60) as response, open(path, "wb") as f:
                    f.write(response.read())
            with open(path, "rb") as f:
                image = base64.b64encode(f.read()).decode()
            requests.append((f"image/{name}", {"state": "Assess this image for use in a Wikipedia article.",
                                               "images": [image], "questions": IMAGE_QUESTIONS}))
    if args.requests:
        with open(args.requests) as f:
            for number, line in enumerate(f, 1):
                if line.strip():
                    body = json.loads(line)
                    requests.append((f"file/{body.pop('id', number)}", body))
    return requests


def run(args):
    requests = build_requests(args)
    results = {}
    for number, (rid, body) in enumerate(requests, 1):
        try:
            request = urllib.request.Request(args.url, json.dumps(body).encode(), {"Content-Type": "application/json"})
            with LOCAL.open(request, timeout=args.timeout) as response:
                data = json.load(response)
            results[rid] = {"answers": data["answers"], "input_tokens": data["usage"]["input_tokens"]}
        except Exception as e:  # noqa: BLE001 - an error is a result, not a crash
            results[rid] = {"error": repr(e)[:300]}
        print(f"[{number}/{len(requests)}] {'ERROR' if 'error' in results[rid] else 'ok   '} {rid}", flush=True)
    with open(args.out, "w") as f:
        json.dump({"label": args.label, "url": args.url, "results": results}, f, ensure_ascii=False, indent=1)
    print(f"Saved {len(results)} results to {args.out}")


def probabilities(answer):
    """Return the probability of each option of one answer."""
    if answer["type"] == "noul":
        return {"yes": answer["noul"]}
    return answer["probabilities"]


def decision(answer):
    """Return the answer as a person reads it, and its probability.

    That is yes/no at 0.5, the choice, or the most likely score level.
    """
    if answer["type"] == "noul":
        p = answer["noul"]
        return ("yes", p) if p >= 0.5 else ("no", 1 - p)
    if answer["type"] == "choice":
        return answer["choice"], answer["probabilities"][answer["choice"]]
    level = max(answer["probabilities"], key=answer["probabilities"].get)
    return level, answer["probabilities"][level]


def parity_lines(path_a, path_b, max_change_limit, min_agreement, sub="##", verdict=True):
    """Compare two saved runs. Return the Markdown lines and whether they pass.

    With verdict=False, the lines give the numbers without PASS/FAIL or the limits.
    """
    with open(path_a) as f:
        a = json.load(f)
    with open(path_b) as f:
        b = json.load(f)
    la, lb = a["label"], b["label"]
    changes, flips, errors, token_mismatch = [], [], [], []
    same = {"text": [0, 0], "image": [0, 0], "file": [0, 0]}
    for rid in sorted(set(a["results"]) & set(b["results"])):
        ra, rb = a["results"][rid], b["results"][rid]
        if "error" in ra or "error" in rb:
            errors.append(rid)
            continue
        if ra["input_tokens"] != rb["input_tokens"]:
            token_mismatch.append(rid)
        kind = rid.split("/")[0]
        for qid, qa in ra["answers"].items():
            qb = rb["answers"][qid]
            pa, pb = probabilities(qa), probabilities(qb)
            for option in pa:
                changes.append((abs(pa[option] - pb[option]), kind, rid, qid, option, pa[option], pb[option]))
            same[kind][1] += 1
            (da, pda), (db, pdb) = decision(qa), decision(qb)
            if da == db:
                same[kind][0] += 1
            else:
                flips.append((rid, qid, da, pda, db, pdb))

    def stats(kind=None):
        values = sorted(c[0] for c in changes if kind in (None, c[1]))
        if not values:
            return None
        p99 = values[max(0, int(0.99 * len(values)) - 1)]
        return statistics.mean(values), p99, values[-1]

    total_same = sum(s for s, _ in same.values())
    total = sum(t for _, t in same.values())
    max_change = stats()[2] if changes else float("nan")
    agreement = total_same / total if total else 0.0
    passed = not errors and not token_mismatch and max_change <= max_change_limit and agreement >= min_agreement

    if verdict:
        summary = (
            f"**{'PASS' if passed else 'FAIL'}**: largest change {max_change:.4f} (limit {max_change_limit}), "
            f"same answer {agreement:.1%} (limit {min_agreement:.0%}), "
        )
    else:
        summary = f"Largest change {max_change:.4f}, same answer {agreement:.1%}, "
    lines = [
        summary + f"{len(errors)} failed requests, {len(token_mismatch)} with different token counts.",
        "",
        "| Requests | Questions | Mean change | 99th percentile | Largest change | Same answer |",
        "|---|---|---|---|---|---|",
    ]
    for kind in ("text", "image", "file"):
        if same[kind][1]:
            mean, p99, top = stats(kind)
            lines.append(f"| {kind} | {same[kind][1]} | {mean:.4f} | {p99:.4f} | {top:.4f} | "
                         f"{same[kind][0]}/{same[kind][1]} |")
    lines += [
        "",
        f"*Change* is the difference between the {la} and {lb} probability of one option (0 to 1). "
        "*Same answer* means the same choice, the same yes/no at 0.5, or the same most likely score level.",
        "",
        f"{sub} Largest changes",
        "",
        f"| Request | Question | Option | {la} | {lb} | Change |",
        "|---|---|---|---|---|---|",
    ]
    for change, _, rid, qid, option, pa, pb in sorted(changes, reverse=True)[:10]:
        lines.append(f"| {rid} | {qid} | {option} | {pa:.4f} | {pb:.4f} | {change:.4f} |")
    if flips:
        lines += ["", f"{sub} Answers that changed", "", f"| Request | Question | {la} | {lb} |", "|---|---|---|---|"]
        for rid, qid, da, pa, db, pb in flips:
            lines.append(f"| {rid} | {qid} | {da} ({pa:.2f}) | {db} ({pb:.2f}) |")
    if errors or token_mismatch:
        lines += ["", f"{sub} Problems", ""]
        lines += [f"- Failed request: {rid}" for rid in errors]
        lines += [f"- Different token count: {rid}" for rid in token_mismatch]
    return lines, passed


def write_report(path, lines, passed):
    text = "\n".join(lines) + "\n"
    with open(path, "w") as f:
        f.write(text)
    print(text)
    print(f"Wrote {path}")
    sys.exit(0 if passed else 1)


def compare(args):
    lines, passed = parity_lines(args.a, args.b, args.max_change, args.min_agreement)
    with open(args.b) as f:
        lb = json.load(f)["label"]
    with open(args.a) as f:
        la = json.load(f)["label"]
    write_report(args.report, [f"# Clef parity: {lb} against {la}", ""] + lines, passed)


def server_facts(path):
    """Read what vLLM logged about quantization and GPU memory."""
    facts = {"quantization": "?", "weights": "?", "kv_cache": "?"}
    if not os.path.exists(path):
        return facts
    with open(path, errors="replace") as f:
        log = f.read()
    for key, pattern in (
        ("quantization", r"quantization=([\w-]+)"),
        ("weights", r"Model loading took ([\d.]+ GiB)"),
        ("kv_cache", r"Available KV cache memory: ([\d.]+ GiB)"),
    ):
        match = re.search(pattern, log)
        if match:
            facts[key] = match.group(1)
    if facts["quantization"] == "None":
        facts["quantization"] = "none (BF16)"
    return facts


def report(args):
    """Write one report for the whole run: quantization, GPU memory and answers."""
    d = args.results
    parity, passed = parity_lines(os.path.join(d, "bf16.json"), os.path.join(d, "fp8.json"),
                                  args.max_change, args.min_agreement, sub="###", verdict=not args.no_verdict)
    when = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    lines = [
        "# Clef-Flash: FP8 against BF16",
        "",
        (f"**{'PASS' if passed else 'FAIL'}**. " if not args.no_verdict else "") + f"Run on {socket.gethostname()} on {when}.",
        *([""] + args.note if args.note else []),
        "",
    ]
    manifest_path = os.path.join(d, "MANIFEST.json")
    if os.path.exists(manifest_path):
        with open(manifest_path) as f:
            m = json.load(f)
        v = m["versions"]
        lines += [
            "## Quantization",
            "",
            f"- Source: `{m['source']}`. Every input file matched Cloudflare's published hash.",
            f"- Tool: llm-compressor {v['llmcompressor']}, scheme `{m['recipe']['scheme']}` "
            f"(compressed-tensors {v['compressed_tensors']}, transformers {v['transformers']}, torch {v['torch']}).",
            f"- {m['fp8_weights']} weights in FP8, relative error mean {m['relative_error_mean']:.2%}, "
            f"max {m['relative_error_max']:.2%}. All other tensors are bit-identical to Cloudflare's.",
            "",
        ]
    bf16, fp8 = server_facts(os.path.join(d, "server-bf16.log")), server_facts(os.path.join(d, "server-fp8.log"))
    lines += [
        "## GPU memory, as logged by vLLM",
        "",
        "| | BF16 | FP8 |",
        "|---|---|---|",
        f"| Quantization | {bf16['quantization']} | {fp8['quantization']} |",
        f"| Weights | {bf16['weights']} | {fp8['weights']} |",
        f"| Free for KV cache | {bf16['kv_cache']} | {fp8['kv_cache']} |",
        "",
        "Speed on the MI210 means nothing here: it has no FP8 hardware, so vLLM emulates FP8.",
        "",
        "## Answers",
        "",
        *parity,
    ]
    write_report(args.out or os.path.join(d, "report.md"), lines, passed)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)
    p = sub.add_parser("run", help="Send the requests to one server and save the answers.")
    p.add_argument("--label", required=True, help="Name of this server, for example bf16 or fp8.")
    p.add_argument("--out", required=True)
    p.add_argument("--url", default=URL)
    p.add_argument("--requests", help="Extra requests: a JSONL file, one request body per line.")
    p.add_argument("--image-dir", default="parity-images")
    p.add_argument("--no-images", action="store_true")
    p.add_argument("--timeout", type=float, default=300)
    p.set_defaults(func=run)
    c = sub.add_parser("compare", help="Compare two saved runs and write the report.")
    c.add_argument("a", help="Reference run, for example bf16.json.")
    c.add_argument("b", help="Candidate run, for example fp8.json.")
    c.add_argument("--report", default="parity-report.md")
    c.add_argument("--max-change", type=float, default=0.03)
    c.add_argument("--min-agreement", type=float, default=0.99)
    c.set_defaults(func=compare)
    r = sub.add_parser("report", help="Write one report for a run_all.sh results directory.")
    r.add_argument("results", help="Directory with bf16.json, fp8.json, MANIFEST.json and the server logs.")
    r.add_argument("--note", action="append", help="A line to add under the title. Repeat for more lines.")
    r.add_argument("--no-verdict", action="store_true", help="Leave out PASS/FAIL and the limits, for example to publish.")
    r.add_argument("--out", help="Where to write the report. Default: report.md in the results directory.")
    r.add_argument("--max-change", type=float, default=0.03)
    r.add_argument("--min-agreement", type=float, default=0.99)
    r.set_defaults(func=report)
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()