fund_rfid_data/score_baseline.py

#!/usr/bin/env python3
"""
Baseline scorer for the SEC prospectus -> RDF triple PoC.

Compares a system's predicted triples against the N-CEN gold graph and reports
triple-level precision / recall / F1 (the metric used in the thesis).

Two intended uses:
  1. NON-MODEL baseline -- the gold itself is from N-CEN; this script also runs a
     trivial string-match baseline that scans the prospectus prose for each gold
     entity label and emits the edge if the object's name appears in the text.
     This gives a "no-LLM" lower bound: how many edges are even surfaceable by
     naive string matching, and how the prose supports each relation type.
  2. STRONG-MODEL baseline -- point --pred at a JSONL of model predictions
     (one obj per line: {"cik": ..., "triples": [{"s","p","o"}, ...]}) to score
     GPT-4 / Opus extractions against the same gold.

Triple matching is on (subject_type, predicate, normalized_object_label) so that
IRI slug differences do not cause false negatives; object org names are normalized
(lowercased, legal suffixes stripped) before comparison.

Usage:
  python score_baseline.py stringmatch            # no-model baseline over the PoC samples
  python score_baseline.py model --pred preds.jsonl
"""

import argparse
import json
import re
from collections import defaultdict
from pathlib import Path

SAMPLES = Path("data/rdf_poc/samples.jsonl")
GOLD = Path("data/rdf_poc/gold_graphs.jsonl")

_SUFFIX = re.compile(
    r"\b(llc|l\.l\.c|inc|inc\.|incorporated|corp|corporation|company|co|co\.|"
    r"ltd|limited|lp|l\.p|llp|na|n\.a|trust|the)\b", re.I)


def norm(name: str) -> str:
    s = (name or "").lower()
    s = _SUFFIX.sub(" ", s)
    s = re.sub(r"[^a-z0-9]+", " ", s)
    return re.sub(r"\s+", " ", s).strip()


def triple_key(t, entities):
    """Match key: (predicate, normalized object label).

    Subject type is intentionally omitted: gold subjects are IRIs resolved via
    `entities`, but model predictions carry raw subject strings, so a type-based
    key would never match predictions. Within one fund's scope, (predicate,
    object name) uniquely identifies a service-provider edge. The object label is
    resolved through `entities` for gold IRIs and falls back to the raw string for
    predictions.
    """
    o_label = entities.get(t["o"], {}).get("label", t["o"])
    return (t["p"], norm(o_label))


# ---- fuzzy entity-name matching --------------------------------------------
# The same entity is named differently in the prospectus text (what a model
# copies) and in N-CEN gold, e.g. "UBS AM (Americas)" vs
# "UBS Asset Management (Americas) LLC". Strict equality undercounts true hits,
# so fuzzy matching treats two names as the same entity when their token sets
# overlap strongly or one is a token-subset of the other.
_STOP = {"of", "and", "the", "fund", "services", "service", "asset", "management",
         "investment", "investments", "company", "co"}


def _toks(name_norm):
    return {w for w in name_norm.split() if w not in _STOP and len(w) > 1}


def names_match(a_norm, b_norm, thresh=0.6):
    if a_norm == b_norm:
        return True
    if not a_norm or not b_norm:
        return False
    if a_norm in b_norm or b_norm in a_norm:   # one is a substring of the other
        return True
    ta, tb = _toks(a_norm), _toks(b_norm)
    if not ta or not tb:
        return False
    if ta <= tb or tb <= ta:                   # token-subset (after stopwords)
        return True
    jacc = len(ta & tb) / len(ta | tb)
    return jacc >= thresh


def fuzzy_tp(gold_keys, pred_keys):
    """Greedy count of gold edges matched by some prediction (fuzzy on name).

    gold_keys / pred_keys are lists of (predicate, normalized_name). A gold edge
    matches a prediction if predicates are equal and names_match() holds; each
    prediction is consumed at most once.
    """
    used = [False] * len(pred_keys)
    tp = 0
    for gp, gn in gold_keys:
        for i, (pp, pn) in enumerate(pred_keys):
            if not used[i] and pp == gp and names_match(gn, pn):
                used[i] = True
                tp += 1
                break
    return tp


def score(gold_triples, pred_triples, entities):
    g = {triple_key(t, entities) for t in gold_triples}
    p = {triple_key(t, entities) for t in pred_triples}
    tp = len(g & p)
    prec = tp / len(p) if p else 0.0
    rec = tp / len(g) if g else 0.0
    f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0
    return tp, len(p), len(g), prec, rec, f1


def stringmatch_predict(sample):
    """No-model baseline: emit a gold edge iff the object label appears in the prose."""
    text_norm = norm(sample["input_text"])
    ents = sample.get("entities") or {}
    # samples.jsonl doesn't carry entities; rebuild a minimal label map from gold
    preds = []
    for t in sample["target_triples"]:
        o_label = sample["_entities"].get(t["o"], {}).get("label", "")
        if o_label and norm(o_label) and norm(o_label) in text_norm:
            preds.append(t)
    return preds


def load_samples_with_entities():
    gold = {json.loads(l)["cik"]: json.loads(l) for l in open(GOLD, encoding="utf-8")}
    out = []
    for l in open(SAMPLES, encoding="utf-8"):
        s = json.loads(l)
        g = gold.get(s["cik"], {})
        s["_entities"] = g.get("entities", {})
        out.append(s)
    return out


def run_stringmatch():
    samples = load_samples_with_entities()
    agg = defaultdict(lambda: [0, 0, 0])  # predicate -> [tp, npred, ngold]
    micro = [0, 0, 0]
    print(f"{'CIK':<12}{'trust':<34}{'P':>6}{'R':>6}{'F1':>6}  (string-match, no model)")
    print("-" * 76)
    for s in samples:
        ents = s["_entities"]
        gold_t = s["target_triples"]
        pred_t = stringmatch_predict(s)
        tp, np_, ng, prec, rec, f1 = score(gold_t, pred_t, ents)
        micro[0] += tp; micro[1] += np_; micro[2] += ng
        for t in gold_t:
            agg[t["p"]][2] += 1
        gk = {triple_key(t, ents) for t in gold_t}
        for t in pred_t:
            agg[t["p"]][1] += 1
            if triple_key(t, ents) in gk:
                agg[t["p"]][0] += 1
        print(f"{s['cik']:<12}{s['trust_name'][:32]:<34}{prec:>6.2f}{rec:>6.2f}{f1:>6.2f}")
    tp, np_, ng = micro
    P = tp / np_ if np_ else 0
    R = tp / ng if ng else 0
    F = 2 * P * R / (P + R) if (P + R) else 0
    print("-" * 76)
    print(f"MICRO over {len(samples)} samples:  P={P:.3f}  R={R:.3f}  F1={F:.3f}  (tp={tp}, pred={np_}, gold={ng})")
    print("\nPer-relation recall of the no-model string-match baseline:")
    print(f"  {'relation':<16}{'recall':>8}{'gold':>8}  (how prose-grounded each edge type is)")
    for p, (t, npd, ngd) in sorted(agg.items(), key=lambda x: -x[1][2]):
        r = t / ngd if ngd else 0
        print(f"  {p:<16}{r:>8.2f}{ngd:>8}")


def run_model(pred_path, fuzzy=True):
    samples = load_samples_with_entities()
    # join predictions by sample_id (per-fund) when present, else by cik
    preds_by_id = {}
    preds_by_cik = defaultdict(list)
    for l in open(pred_path, encoding="utf-8"):
        r = json.loads(l)
        if r.get("sample_id"):
            preds_by_id[r["sample_id"]] = r.get("triples", [])
        preds_by_cik[r["cik"]] = r.get("triples", [])

    micro = [0, 0, 0]
    per_rel = defaultdict(lambda: [0, 0, 0])  # predicate -> [tp, npred, ngold]
    for s in samples:
        ents = s["_entities"]
        pred = preds_by_id.get(s.get("sample_id")) if s.get("sample_id") in preds_by_id \
            else preds_by_cik.get(s["cik"], [])
        gold_keys = [triple_key(t, ents) for t in s["target_triples"]]
        pred_keys = [triple_key(t, ents) for t in pred]
        if fuzzy:
            tp_s = fuzzy_tp(gold_keys, pred_keys)
        else:
            gk = set(gold_keys); tp_s = len(gk & set(pred_keys))
        micro[0] += tp_s; micro[1] += len(pred_keys); micro[2] += len(gold_keys)
        # per-relation (restrict matching within the same predicate)
        for rel in set(p for p, _ in gold_keys) | set(p for p, _ in pred_keys):
            g_r = [k for k in gold_keys if k[0] == rel]
            p_r = [k for k in pred_keys if k[0] == rel]
            t_r = fuzzy_tp(g_r, p_r) if fuzzy else len(set(g_r) & set(p_r))
            per_rel[rel][0] += t_r
            per_rel[rel][1] += len(p_r)
            per_rel[rel][2] += len(g_r)

    tp, np_, ng = micro
    P = tp / np_ if np_ else 0
    R = tp / ng if ng else 0
    F = 2 * P * R / (P + R) if (P + R) else 0
    mode = "fuzzy" if fuzzy else "strict"
    print(f"\nMODEL baseline over {len(samples)} samples ({mode} name matching):")
    print(f"  MICRO  P={P:.3f}  R={R:.3f}  F1={F:.3f}  (tp={tp}, pred={np_}, gold={ng})")
    print(f"\n  {'relation':16s}{'P':>6}{'R':>6}{'F1':>6}{'gold':>7}")
    for p, (t, npd, ngd) in sorted(per_rel.items(), key=lambda x: -x[1][2]):
        pr = t / npd if npd else 0
        rc = t / ngd if ngd else 0
        f1 = 2 * pr * rc / (pr + rc) if (pr + rc) else 0
        print(f"  {p:16s}{pr:>6.2f}{rc:>6.2f}{f1:>6.2f}{ngd:>7}")


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd")
    sub.add_parser("stringmatch")
    m = sub.add_parser("model"); m.add_argument("--pred", required=True)
    m.add_argument("--strict", action="store_true",
                   help="exact normalized name matching (default: fuzzy)")
    args = ap.parse_args()
    if args.cmd == "stringmatch":
        run_stringmatch()
    elif args.cmd == "model":
        run_model(args.pred, fuzzy=not args.strict)
    else:
        ap.print_help()


if __name__ == "__main__":
    main()