bundle/scripts/claim5_judge.py

#!/usr/bin/env python
"""
Judge Claim 5 completions with a local LLM (open substitute for GPT-4o-mini).
Concept score 0-2 (alignment with target concept) and fluency 0-2 (coherence).
Final steering score per (method) = harmonic-mean(concept, fluency), best alpha per
concept, then median across concepts. Mirrors AxBench aggregation used by the paper.
"""
import os, sys, json, argparse, re, statistics
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

JUDGE = "meta-llama/Llama-3.2-3B-Instruct"

def build_prompt(concept, text):
    return (f"You are rating a text sample.\n"
            f"Target concept: {concept}\n"
            f"Text: \"{text}\"\n\n"
            f"Rate two things, each 0-2:\n"
            f"- concept: 2 = clearly about the target concept, 1 = loosely related, 0 = unrelated.\n"
            f"- fluency: 2 = coherent English, 1 = somewhat broken, 0 = gibberish/repetition.\n"
            f"Answer EXACTLY as: concept=<0-2> fluency=<0-2>")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--inp", default="outputs/claim5_completions.json")
    ap.add_argument("--out", default="outputs/claim5_scores.json")
    args = ap.parse_args()

    data = json.load(open(args.inp))
    tokj = AutoTokenizer.from_pretrained(JUDGE)
    model = AutoModelForCausalLM.from_pretrained(JUDGE, torch_dtype=torch.bfloat16, device_map="cuda").eval()

    @torch.no_grad()
    def score(concept, text):
        msg = [{"role": "user", "content": build_prompt(concept, text)}]
        enc = tokj.apply_chat_template(msg, add_generation_prompt=True, return_tensors="pt", return_dict=True)
        enc = {k: v.to("cuda") for k, v in enc.items()}
        n_in = enc["input_ids"].shape[1]
        out = model.generate(**enc, max_new_tokens=16, do_sample=False, pad_token_id=tokj.eos_token_id)
        txt = tokj.decode(out[0, n_in:], skip_special_tokens=True)
        c = re.search(r"concept\s*=\s*([0-2])", txt); f = re.search(r"fluency\s*=\s*([0-2])", txt)
        return (int(c.group(1)) if c else 0, int(f.group(1)) if f else 0)

    def hmean(c, f):
        c2, f2 = c/2.0, f/2.0
        return 0.0 if (c2+f2) == 0 else 2*c2*f2/(c2+f2)

    def score_set(concept, comps):
        vals = []
        for t in comps:
            gen = t.split("I think that", 1)[-1][:300]
            c, f = score(concept, gen)
            vals.append(hmean(c, f))
        return statistics.mean(vals) if vals else 0.0

    results = {"judge": JUDGE, "per_concept": [], "mfa_best": [], "diffmeans_best": []}
    base_scores = []
    for c in data["concepts"]:
        concept = ", ".join(c["concept_tokens"])
        mfa_by_a = {a: score_set(concept, comps) for a, comps in c["mfa"].items()}
        dm_by_a = {a: score_set(concept, comps) for a, comps in c.get("diffmeans", {}).items()}
        mfa_best = max(mfa_by_a.values()) if mfa_by_a else 0.0
        dm_best = max(dm_by_a.values()) if dm_by_a else 0.0
        results["per_concept"].append({"comp": c["comp"], "kind": c["kind"], "concept": concept,
                                       "mfa_best": round(mfa_best, 3), "diffmeans_best": round(dm_best, 3),
                                       "mfa_by_alpha": {k: round(v,3) for k,v in mfa_by_a.items()},
                                       "diffmeans_by_alpha": {k: round(v,3) for k,v in dm_by_a.items()}})
        results["mfa_best"].append(mfa_best); results["diffmeans_best"].append(dm_best)
        print(f"comp {c['comp']:5d} [{c['kind']:6}] {concept[:40]:40} MFA={mfa_best:.3f} DiffMeans={dm_best:.3f}")

    # baseline concept-agnostic fluency proxy (concept=any -> just fluency)
    results["summary"] = {
        "n_concepts": len(results["mfa_best"]),
        "mfa_median": round(statistics.median(results["mfa_best"]), 3),
        "diffmeans_median": round(statistics.median(results["diffmeans_best"]), 3),
        "mfa_mean": round(statistics.mean(results["mfa_best"]), 3),
        "diffmeans_mean": round(statistics.mean(results["diffmeans_best"]), 3),
    }
    s = results["summary"]
    s["mfa_over_diffmeans_ratio_median"] = round(s["mfa_median"]/max(s["diffmeans_median"],1e-6), 2)
    json.dump(results, open(args.out, "w"), indent=2)
    print(json.dumps(results["summary"], indent=2))


if __name__ == "__main__":
    main()