Blame for sim.py

"""Replay the stream once and write replay.json, which the app plays back.

For each request in order: if a rule set exists and the gate passes, the trusted, not switched-off rules answer;
otherwise the teacher answers and its answer becomes a labelled example. SPOT_EVERY-th rule answers (by request-id
hash) are also sent to the teacher: that is the drift guard's evidence, and those answers are labelled too. A rule
whose spot-check precision falls below GUARD_PREC after GUARD_MIN checked spans is switched off for good. Each time the
labelled set reaches the next size in SCHEDULE, RuleChef's search synthesis (no LLM) re-learns the rules from all
labelled examples; the new set takes over at the next request. Every rule set version is also scored on the
fixed holdout against the teacher and against the TAB annotation.
Teacher answers are read from teacher.jsonl (precomputed with run_teacher.py, the same code the app runs live).
usage: sim.py DATA_DIR OUT.json
"""
import hashlib
import json
import os
import platform
import sys
import time

import rulechef
from rulechef.core import Dataset, Example, Task, TaskType
from rulechef.search import SearchConfig, SearchSynthesizer

import cascade as C

SPOT_EVERY = 10
GUARD_MIN, GUARD_PREC = int(os.environ.get("GUARD_MIN", 5)), float(os.environ.get("GUARD_PREC", 0.8))
SCHEDULE = [50, 100] + list(range(200, 5001, 100))
MIN_SUPPORT = 2


def spot(rid):
    return int(hashlib.sha1(rid.encode()).hexdigest(), 16) % SPOT_EVERY == 0


def dataset(labelled):
    task = Task(name="stream", description="date / code / quantity spans", input_schema={"text": "str"},
                output_schema={"entities": "List[Entity]"}, type=TaskType.NER, text_field="text")
    ds = Dataset(name="stream", task=task)
    for i, (text, spans) in enumerate(labelled):
        ents = [{"text": s["text"], "start": s["start"], "end": s["end"], "type": s["type"]} for s in spans]
        ds.examples.append(Example(id=str(i), input={"text": text}, expected_output={"entities": ents}, source="teacher"))
    return ds


def rule_key(d):
    return d["output_template"]["type"] + "\t" + d["content"] + "\t" + d["output_template"].get("text", "$0")


def rule_stats(rules, labelled):
    """Each rule alone on the labelled data, against the teacher: matches and precision."""
    out = {}
    for r in rules:
        tp = fp = 0
        for text, spans in labelled:
            t, f, _ = C.match_counts(C.run_rules([r], text), spans)
            tp, fp = tp + t, fp + f
        out[r.id] = (tp, fp)
    return out


def evaluate(holdout, rules, gate):
    """Holdout scores for one rule set (rules = the trusted, active ones)."""
    agg = {k: [0, 0, 0] for k in ("rules_t", "casc_t", "rules_g", "casc_g", "teach_g", "routed_t")}
    passed = exact = 0
    for r in holdout:
        rs = C.run_rules(rules, r["text"]) if gate is not None else []
        ok = gate is not None and not gate.unexplained(r["text"], rs)
        casc = rs if ok else r["teacher"]
        for k, pred, ref in (("rules_t", rs, r["teacher"]), ("casc_t", casc, r["teacher"]), ("rules_g", rs, r["gold"]),
                             ("casc_g", casc, r["gold"]), ("teach_g", r["teacher"], r["gold"])):
            for j, v in enumerate(C.match_counts(pred, ref)):
                agg[k][j] += v
        if ok:
            passed += 1
            c = C.match_counts(rs, r["teacher"])
            exact += c[1] == 0 and c[2] == 0
            for j, v in enumerate(c):
                agg["routed_t"][j] += v
    res = {"n": len(holdout), "rules_answer": passed, "exact_agree": exact}
    for k, v in agg.items():
        p, rr, f = C.prf(*v)
        res[k] = {"tp": v[0], "fp": v[1], "fn": v[2], "P": round(p, 4), "R": round(rr, 4), "F1": round(f, 4)}
    return res


def main(data, out):
    reqs = [json.loads(line) for line in open(f"{data}/requests.jsonl")]
    teach = {}
    for line in open(f"{data}/teacher.jsonl"):
        t = json.loads(line)
        teach[t["id"]] = t
    for r in reqs:
        r["teacher"] = teach[r["id"]]["spans"]
        r["teacher_ms"] = teach[r["id"]]["ms"]
    stream = [r for r in reqs if r["part"] == "stream"]
    holdout = [r for r in reqs if r["part"] == "holdout"]

    labelled, versions, events, guard_log = [], [], [], []
    guard, blocked = {}, set()
    rules, rule_dicts, trusted, gate, next_i = [], {}, set(), None, 0
    versions.append({"v": 0, "at": 0, "n_labelled": 0, "synth_s": 0.0, "rules": [], "gate": None,
                     "holdout": evaluate(holdout, [], None)})

    def active():
        return [r for r in rules if r.id in trusted and rule_key(rule_dicts[r.id]) not in blocked]

    for i, r in enumerate(stream):
        if next_i < len(SCHEDULE) and len(labelled) >= SCHEDULE[next_i]:
            while next_i < len(SCHEDULE) and len(labelled) >= SCHEDULE[next_i]:
                next_i += 1
            t0 = time.perf_counter()
            syn = SearchSynthesizer(SearchConfig(labels=C.TYPES, min_support=MIN_SUPPORT))
            rules = syn.synthesize(dataset(labelled))
            synth_s = time.perf_counter() - t0
            t0 = time.perf_counter()
            gate = C.Gate.learn(labelled)
            stats = rule_stats(rules, labelled)
            rule_dicts = {x.id: x.to_dict() for x in rules}
            trusted = {x.id for x in rules if sum(stats[x.id]) >= C.TRUST_MIN and stats[x.id][0] / sum(stats[x.id]) >= C.TRUST_PREC}
            post_s = time.perf_counter() - t0
            v = len(versions)
            versions.append({
                "v": v, "at": i, "n_labelled": len(labelled), "synth_s": round(synth_s, 2), "check_s": round(post_s, 2),
                "rules": [{"id": x.id, "key": rule_key(rule_dicts[x.id]), "type": rule_dicts[x.id]["output_template"]["type"],
                           "regex": x.content, "template": rule_dicts[x.id]["output_template"].get("text", "$0"),
                           "tp": stats[x.id][0], "fp": stats[x.id][1], "trusted": x.id in trusted, "dict": rule_dicts[x.id]}
                          for x in rules],
                "gate": gate.to_dict(), "blocked_at_deploy": sorted(blocked), "report": syn.report,
                "holdout": evaluate(holdout, active(), gate)})
            print(f"v{v} at {i} labelled {len(labelled)} rules {len(rules)} trusted {len(trusted)} synth {synth_s:.1f}s "
                  f"holdout rules-answer {versions[-1]['holdout']['rules_answer']} casc F1 {versions[-1]['holdout']['casc_t']['F1']}", flush=True)

        t0 = time.perf_counter()
        act = active()
        rs = C.run_rules(act, r["text"]) if gate is not None else []
        unexpl = gate.unexplained(r["text"], rs) if gate is not None else ["(no rules yet)"]
        rule_ms = (time.perf_counter() - t0) * 1000
        ev = {"id": r["id"], "v": len(versions) - 1, "rule_ms": round(rule_ms, 3), "teacher_ms": r["teacher_ms"],
              "fired": sorted({rule_key(rule_dicts[s["rule_id"]]) for s in rs}) if not unexpl else []}
        if unexpl:
            ev["by"] = "teacher"
            labelled.append((r["text"], r["teacher"]))
        else:
            ev["by"] = "rules"
            if spot(r["id"]):
                ev["checked"] = True
                tp, fp, fn = C.match_counts(rs, r["teacher"])
                ev["agree"] = fp == 0 and fn == 0
                labelled.append((r["text"], r["teacher"]))
                tkeys = [C.key(s) for s in r["teacher"]]
                for s in rs:
                    k = rule_key(rule_dicts[s["rule_id"]])
                    g = guard.setdefault(k, [0, 0])
                    hit = C.key(s) in tkeys
                    if hit:
                        tkeys.remove(C.key(s))
                    g[0 if hit else 1] += 1
                    if k not in blocked and sum(g) >= GUARD_MIN and g[0] / sum(g) < GUARD_PREC:
                        blocked.add(k)
                        guard_log.append({"at": i, "key": k, "tp": g[0], "fp": g[1]})
                        print("guard: switched off", k.replace("\t", "  "), g, flush=True)
        events.append(ev)

    meta = {"n_stream": len(stream), "n_holdout": len(holdout), "spot_every": SPOT_EVERY, "guard_min": GUARD_MIN,
            "guard_prec": GUARD_PREC, "schedule": SCHEDULE, "min_support": MIN_SUPPORT, "trust_min": C.TRUST_MIN,
            "trust_prec": C.TRUST_PREC, "word_min": C.WORD_MIN, "word_ratio": C.WORD_RATIO, "skip_min": C.SKIP_MIN,
            "skip_ratio": C.SKIP_RATIO, "rulechef": getattr(rulechef, "__version__", "?"), "python": platform.python_version(),
            "cpu": platform.processor() or platform.machine()}
    json.dump({"meta": meta, "versions": versions, "events": events, "guard": guard_log, "guard_final": guard},
              open(out, "w"), separators=(",", ":"))
    n_rules = sum(e["by"] == "rules" for e in events)
    calls = sum(e["by"] == "teacher" or e.get("checked", False) for e in events)
    print("rules answered", n_rules, "of", len(events), "model calls", calls, "versions", len(versions) - 1, "switched off", len(guard_log))


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])