esh-pfi-infrastructure/scripts/r49-corpus/entities.py

"""R49 Stage D2 — build the per-work entity map, deterministically.

Technique is F02's, which took three generations to get right and whose lesson is
one of level rather than of cleverness: **the entity map is built once per work,
so the detector must see the work, not the paragraph.**

  v1 position-based   -> MISSES names that start sentences (characters do, constantly)
  v2 dictionary-based -> MISSES names that are words (fiction names people after flowers)
  v3 corpus cap-ratio -> works. No wordlist, no position rule, no LLM.

A token's capitalised count against its lowercase count across the WHOLE work
separates `Jane` (only ever capitalised) from `Door` (capitalised only when it
starts a sentence). Identity linking then joins adjacent capitalised pairs that
recur, which is also what recovers the first-person narrator's gender -- her name
appears mainly in dialogue, surrounded by other people's pronouns, so proximity
inference is structurally blind to exactly the character the adapter is being
trained on.

Nothing here guesses. Unresolved entities block corpus emission and go to a human
pass: held is cheap, wrong is poison -- a silently mis-gendered entity scrambles
pronoun agreement through every renamed copy and nothing downstream would catch it.

⚠ v4, added for BabyYarros: a MID-SENTENCE test on top of the ratio.
The cap/lowercase ratio calls `Hey`, `Holy`, `Hopefully`, `Yep`, `Whoa`, `Nope`
and `Ugh` names, because a dialogue-heavy contemporary novel opens sentences with
them constantly and never writes them lowercase. The v1 lesson was that POSITION
ALONE misses names that start sentences; position as a SECOND filter has no such
problem, because a real name also appears mid-sentence. Measured on BabyYarros the
two populations do not overlap: 33 verified names sit at 0.567-0.985 mid-sentence,
and 19 verified interjections at 0.000-0.222. The gap is 2.5x wide, so the
threshold is not a tuned parameter.

It is OPT-IN (`--min-mid-ratio`, default 0 = off) so the Brontë run stays
byte-reproducible. A 19th-century novel does not have this failure mode in the
same volume, and an unmeasured change to a settled corpus is not an improvement.
"""
from __future__ import annotations
import argparse, collections, json, re, sys
from pathlib import Path

WORD = re.compile(r"\b[A-Za-zÀ-ÿŒœÆæ][a-zà-ÿœæ'’\-]*\b")
TOKEN = re.compile(r"[A-Za-zÀ-ÿŒœÆæ][A-Za-zà-ÿœæ'’\-]*")

#: Ranks, honorifics and address forms are not names. F02 lost `Colonel Aetos`
#: and `Professor Kaori` to this -- without the stoplist the rename replaces the
#: rank. Kinship terms likewise: `Mom` renamed to `Ingrid` was a v2 defect.
STOP_TITLES = {
    "Mr", "Mrs", "Miss", "Ms", "Dr", "Sir", "Lady", "Lord", "Madam", "Madame",
    "Mademoiselle", "Monsieur", "Master", "Captain", "Colonel", "Major", "General",
    "Professor", "Reverend", "Rev", "Doctor", "Saint", "St", "Aunt", "Uncle",
    "Mother", "Father", "Papa", "Mamma", "Mama", "Brother", "Sister", "Cousin",
    "Grandmother", "Grandfather", "Nurse", "King", "Queen", "Prince", "Princess",
    "Duke", "Duchess", "Earl", "Count", "Countess", "Baron", "Squire", "Parson",
    "Monseigneur", "Mlle", "Mme", "M", "Messrs",
    # modern ranks and address forms, added for BabyYarros
    "Sergeant", "Sgt", "Lieutenant", "Lt", "Corporal", "Admiral", "Commander",
    "Cadet", "Officer", "Agent", "Coach", "Senator", "Majesty", "Highness",
}
#: Days, months, and the language/nation adjectives a 19th-century novel is full
#: of. All are always-capitalised and would otherwise pass the ratio test.
STOP_COMMON = {
    "Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Sunday",
    "January","February","March","April","May","June","July","August",
    "September","October","November","December",
    "English","France","French","England","Britain","British","Europe","European",
    "German","Germany","Belgian","Belgium","Scotch","Scottish","Scotland","Irish",
    "Ireland","Welsh","Wales","Latin","Greek","Italian","Italy","Spanish","Spain",
    "Swiss","Switzerland","Dutch","Holland","Roman","Rome","Catholic","Protestant",
    "Christian","Christ","God","Lord","Heaven","Providence","Bible","Sabbath",
    "Christmas","Easter","London","Paris","Brussels","Yorkshire","I","O","Oh","Ah",
    "Yes","No","Well","Now","Then","But","And","The","A","An","He","She","It","They",
    "You","We","His","Her","My","Your","Their","This","That","There","Here","What",
    "Who","When","Where","Why","How","If","So","As","At","In","On","To","For","Of",
    "Nay","Alas","Madam","Sir","Mademoiselle","Monsieur",
}
#: Structural words from the book's own apparatus. `Chapter` and `Article` pass
#: both the ratio test and the mid-sentence test -- `BONUS CONTENT Chapter Nine`
#: and `Article Three` put them mid-sentence -- and renaming them would rewrite
#: the corpus's own scaffolding.
STOP_STRUCTURAL = {
    "Chapter", "Chapters", "Prologue", "Epilogue", "Part", "Appendix", "Volume",
    "Article", "Section", "Contents", "Content", "Bonus", "Preface", "Interlude",
}
STOP = STOP_TITLES | STOP_COMMON | STOP_STRUCTURAL

MALE_PRON = {"he", "him", "his", "himself"}
FEM_PRON = {"she", "her", "hers", "herself"}


def load(corpus: Path) -> dict[str, str]:
    man = json.loads((corpus / "manifest.json").read_text())
    out = {}
    for w in man["works"]:
        rows = [json.loads(l) for l in (corpus / w["path"]).read_text(encoding="utf-8").splitlines()]
        out[w["slug"]] = "\n\n".join(r["text"] for r in rows)
    return out


#: `’s` is a possessive and the rest are contractions; none of them is part of the
#: name. TOKEN keeps the apostrophe, so without folding `Afendra’s` is its own key.
CLITIC = re.compile(r"[’'](?:s|d|ll|ve|re|m|t)$", re.I)


def detect(text: str, min_count: int, max_ratio: float, fold_clitics: bool = False) -> dict[str, dict]:
    """Corpus-level capitalised-vs-lowercase ratio. See module docstring.

    ⚠ `fold_clitics` folds `Afendra’s` into `Afendra`. Without it an entity that
    NEVER appears unpossessed is keyed with the apostrophe, and both rename.py and
    the leak gate skip apostrophe keys as contractions -- so it is never renamed
    AND never reported. Measured on BabyYarros: `Afendra` survived every copy
    while the gate read 0 of 314, which is the worst failure shape there is.
    """
    cap, low = collections.Counter(), collections.Counter()
    for m in TOKEN.finditer(text):
        t = m.group(0)
        if fold_clitics:
            t = CLITIC.sub("", t) or t
        (cap if t[:1].isupper() else low)[t.lower()] += 1
    ents = {}
    for key, c in cap.items():
        if c < min_count:
            continue
        l = low[key]
        ratio = l / c
        if ratio > max_ratio:
            continue
        # recover the dominant surface spelling
        ents[key] = {"cap": c, "lower": l, "ratio": round(ratio, 4)}
    return ents


#: Whatever can sit between a sentence terminator and the first word of the next
#: sentence: whitespace, opening quotes, brackets, a dash.
_OPENERS = set(' \t\n\u201c\u201d"\'\u2018\u2019([{\u2014\u2013-*')
_TERM = set('.!?\u2026')


def mid_sentence(text: str, keys: set[str], fold_clitics: bool = False) -> tuple[dict[str, int], dict[str, int]]:
    """(mid, total) capitalised occurrences per key.

    `mid` counts the ones whose preceding non-opener character is not a sentence
    terminator -- i.e. the capital is the writer's choice and not the position's.
    """
    mid, tot = collections.Counter(), collections.Counter()
    for m in TOKEN.finditer(text):
        t = m.group(0)
        if fold_clitics:
            t = CLITIC.sub("", t) or t
        if not t[:1].isupper():
            continue
        k = t.lower()
        if k not in keys:
            continue
        tot[k] += 1
        i = m.start() - 1
        while i >= 0 and text[i] in _OPENERS:
            i -= 1
        if i >= 0 and text[i] not in _TERM:
            mid[k] += 1
    return mid, tot


#: A word carrying one of these in front of it is a name, whatever its position
#: statistics say. This is rename.py's title-first idea used as a RESCUE rather
#: than as a gender signal.
_HONORIFIC = (r"(?:Mr|Mrs|Ms|Miss|Dr|Doctor|Professor|Prof|Colonel|Col|Major|General|Gen|"
              r"Captain|Capt|Lieutenant|Lt|Sergeant|Sgt|Cadet|Sir|Madam|Lady|Lord|King|Queen|"
              r"Officer|Agent|Coach|Senator|Judge|Father|Mother|Aunt|Uncle)")


def rescue_signals(text: str, keys: set[str]) -> dict[str, tuple[int, int]]:
    """key -> (honorific-preceded, possessive) counts.

    ⚠ The mid-sentence filter drops real SURNAMES that are only ever used as
    address -- measured here, `Delgado` 18/64, `Schur` 0/10, `Rhee` 0/8, because
    every occurrence is `“Mr. Delgado,”` opening a line of dialogue. Two signals
    separate those from the interjections the filter is FOR: a title in front,
    and a possessive. Measured on BabyYarros, all 19 verified interjections score
    zero on both, and every wrongly-dropped surname scores on at least one.
    """
    hon, poss = collections.Counter(), collections.Counter()
    for m in re.finditer(_HONORIFIC + r"\.?\s+([A-ZÀ-Þ][A-Za-zà-ÿœæ\-]+)", text):
        k = m.group(1).lower()
        if k in keys:
            hon[k] += 1
    for m in re.finditer(r"\b([A-ZÀ-Þ][A-Za-zà-ÿœæ\-]+)[’\']s\b", text):
        k = m.group(1).lower()
        if k in keys:
            poss[k] += 1
    return {k: (hon[k], poss[k]) for k in keys}


ACRONYM = re.compile(r"[A-Z]{2,}s?$")


def ratio_rejects(text: str, min_count: int, max_ratio: float, fold_clitics: bool) -> dict[str, dict]:
    """Candidates frequent enough to matter that the cap/lowercase ratio threw out.

    ⚠ The ratio assumes consistent typography and BabyYarros breaks that: `wilder`
    renders an in-book news article entirely in lowercase, so `eleanor baxter` and
    `ms. baxter` appear uncapitalised three times against 23 capitalised ones --
    ratio 0.13 against a 0.05 bar, and a real character is silently never renamed.
    """
    cap, low = collections.Counter(), collections.Counter()
    for m in TOKEN.finditer(text):
        t = m.group(0)
        if fold_clitics:
            t = CLITIC.sub("", t) or t
        (cap if t[:1].isupper() else low)[t.lower()] += 1
    return {k: {"cap": c, "lower": low[k], "ratio": round(low[k] / c, 4)}
            for k, c in cap.items() if c >= min_count and low[k] / c > max_ratio}


#: ⚠ DELIBERATELY NARROWER than `_HONORIFIC`. The wide list is safe when both
#: sides must be capitalised; matched case-insensitively it readmitted 143 junk
#: tokens (`the`, `says`, `like`, `up`) because `major`, `general`, `father`,
#: `sir` and `agent` are ordinary words in lowercase prose. These five are never
#: anything but a title, and the lowercase arm additionally REQUIRES the period.
_ABBREV = re.compile(r"\b(?:Mr|Mrs|Ms|Dr|Mister|Miss)\b\.?\s+([A-ZÀ-Þ][A-Za-zà-ÿœæ\-]+)"
                     r"|\b(?:mr|mrs|ms|dr)\.\s+([a-zà-ÿœæ][a-zà-ÿœæ\-]+)")


def honorific_hits(text: str, keys: set[str]) -> dict[str, int]:
    """`Miss Baxter` and `ms. baxter` both count; `I miss you` does not."""
    hits = collections.Counter()
    for m in _ABBREV.finditer(text):
        k = (m.group(1) or m.group(2)).lower()
        if k in keys:
            hits[k] += 1
    return hits


def surface_forms(text: str, keys: set[str], prefer_mixed: bool = False,
                  fold_clitics: bool = False) -> dict[str, str]:
    """Dominant spelling per key.

    ⚠ `prefer_mixed` picks the most common NON-all-caps form when one exists.
    Without it a name that happens to sit inside an all-caps passage -- an
    in-world dispatch here, an inscription in Shirley -- gets `BRAEVICK` as its
    surface, and every rule downstream then reasons about an acronym.
    """
    best = collections.defaultdict(collections.Counter)
    for m in TOKEN.finditer(text):
        t = m.group(0)
        if fold_clitics:
            t = CLITIC.sub("", t) or t
        if t[:1].isupper() and t.lower() in keys:
            best[t.lower()][t] += 1
    out = {}
    for k, c in best.items():
        mixed = [(n, f) for f, n in c.most_common() if not ACRONYM.fullmatch(f)]
        out[k] = (max(mixed)[1] if (prefer_mixed and mixed) else c.most_common(1)[0][0])
    return out


def link_identities(text: str, names: set[str], min_pairs: int) -> list[tuple[str, str]]:
    """Adjacent capitalised pairs that recur are one person.

    This is what makes `Xaden Riorson` a single identity so the bare given name
    maps to the given part and the surname to the surname part, keeping the
    honorific form working. It is also what recovers the POV character's gender.
    """
    pairs = collections.Counter()
    toks = [(m.group(0), m.start()) for m in TOKEN.finditer(text)]
    for i in range(len(toks) - 1):
        a, b = toks[i][0], toks[i + 1][0]
        if toks[i + 1][1] - toks[i][1] > len(a) + 2:
            continue  # not actually adjacent
        if a[:1].isupper() and b[:1].isupper() and a not in STOP and b not in STOP:
            if a.lower() in names and b.lower() in names:
                pairs[(a, b)] += 1
    return [p for p, n in pairs.items() if n >= min_pairs]


def resolve_gender(text: str, names: set[str]) -> dict[str, str]:
    """Same-sentence pronoun co-occurrence. Never guesses; unresolved stays unresolved.

    F02: tightening from a +/-200-char window to same-sentence converted a WRONG
    to a HELD while keeping every correct call. Held is cheap; wrong is poison.
    """
    score = collections.defaultdict(lambda: [0, 0])
    for sent in re.split(r"(?<=[.!?])\s+", text):
        low = {w.lower() for w in TOKEN.findall(sent)}
        m, f = bool(low & MALE_PRON), bool(low & FEM_PRON)
        if m == f:
            continue  # both or neither -> no signal
        for t in TOKEN.findall(sent):
            if t[:1].isupper() and t.lower() in names:
                score[t.lower()][0 if m else 1] += 1
    out = {}
    for k, (mm, ff) in score.items():
        tot = mm + ff
        if tot < 3:
            continue
        if mm / tot >= 0.75:
            out[k] = "m"
        elif ff / tot >= 0.75:
            out[k] = "f"
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("corpus")
    ap.add_argument("--out", default=None)
    ap.add_argument("--stoplist", default=None,
                    help="JSON file whose every list value holds surfaces to exclude; "
                         "per-corpus real-world referents, see stoplist_yarros.json")
    ap.add_argument("--rescue-honorific", type=int, default=0,
                    help="readmit a candidate the cap/lowercase ratio rejected when a title "
                         "precedes it at least this many times (0 = off)")
    ap.add_argument("--fold-clitics", action="store_true",
                    help="count `Afendra’s` as `Afendra` so a possessive-only entity is "
                         "detected at all (it is otherwise silently unrenamed AND ungated)")
    ap.add_argument("--drop-acronyms", action="store_true",
                    help="treat an ALWAYS-all-caps surface as an acronym, not a name "
                         "(RSC/ATV/TV/BMX/VIP), and prefer a mixed-case surface when one exists")
    ap.add_argument("--min-count", type=int, default=5)
    ap.add_argument("--max-ratio", type=float, default=0.05)
    ap.add_argument("--min-pairs", type=int, default=2)
    ap.add_argument("--min-mid-ratio", type=float, default=0.0,
                    help="drop a candidate whose capitals are overwhelmingly sentence-initial "
                         "(0 = off, which reproduces the Bronte run)")
    ap.add_argument("--min-mid", type=int, default=2,
                    help="absolute mid-sentence floor, so a 1-of-2 accident cannot qualify")
    ap.add_argument("--control", default="", help="comma-separated known-true names (positive control)")
    ap.add_argument("--negative-control", default="",
                    help="comma-separated known-NON-names that the filter must DROP")
    a = ap.parse_args()
    corpus = Path(a.corpus)
    works = load(corpus)
    stop = set(STOP)
    if a.stoplist:
        blob = json.loads(Path(a.stoplist).read_text())
        extra = {n for v in blob.values() if isinstance(v, list) for n in v}
        stop |= extra
        print(f"  stoplist {a.stoplist}: +{len(extra)} real-world / generic surfaces")

    controls = [c.strip() for c in a.control.split(",") if c.strip()]
    neg_controls = [c.strip() for c in a.negative_control.split(",") if c.strip()]
    report, failed_control = {}, []
    mid_dropped: dict[str, tuple[int, int]] = {}
    rescued: dict[str, tuple[int, int]] = {}
    ratio_rescued: dict[str, tuple[int, int, int]] = {}
    for slug, text in works.items():
        ents = detect(text, a.min_count, a.max_ratio, a.fold_clitics)
        if a.rescue_honorific:
            rej = ratio_rejects(text, a.min_count, a.max_ratio, a.fold_clitics)
            hh = honorific_hits(text, set(rej))
            back = {k: rej[k] for k, n in hh.items() if n >= a.rescue_honorific}
            for k, v in back.items():
                ents.setdefault(k, v)
                ratio_rescued[k] = (hh[k], v["cap"], v["lower"])
        keys = {k for k in ents if k.capitalize() not in stop and k.title() not in stop}
        keys = {k for k in keys if k not in {s.lower() for s in stop}}
        # ⚠ An all-caps surface is an acronym, not a name: RSC, ATV, TV, BMX, VIP,
        # CTDs. Tested on the DOMINANT surface form, because a name also appears
        # inside an all-caps in-world dispatch and must not be lost to that.
        if a.drop_acronyms:
            forms0 = surface_forms(text, keys, prefer_mixed=True, fold_clitics=a.fold_clitics)
            keys = {k for k in keys if not ACRONYM.fullmatch(forms0.get(k, k))}
        if a.min_mid_ratio > 0:
            mid, tot = mid_sentence(text, keys, a.fold_clitics)
            dropped_here = {k for k in keys
                            if mid[k] < a.min_mid or mid[k] / max(tot[k], 1) < a.min_mid_ratio}
            sig = rescue_signals(text, dropped_here)
            rescued_here = {k for k in dropped_here if sum(sig.get(k, (0, 0))) > 0}
            for k in rescued_here:
                rescued[k] = sig[k]
            dropped_here -= rescued_here
            for k in dropped_here:
                mid_dropped[k] = (mid[k], tot[k])
            keys -= dropped_here
        forms = surface_forms(text, keys, prefer_mixed=a.drop_acronyms, fold_clitics=a.fold_clitics)
        links = link_identities(text, keys, a.min_pairs)
        gender = resolve_gender(text, keys)
        # identity linking propagates gender: a bare surname inherits from its given name
        for g, s in links:
            gl, sl = g.lower(), s.lower()
            if gl in gender and sl not in gender:
                gender[sl] = gender[gl]
            elif sl in gender and gl not in gender:
                gender[gl] = gender[sl]
        report[slug] = {"entities": {k: {**ents[k], "surface": forms.get(k, k),
                                         "gender": gender.get(k)} for k in sorted(keys)},
                        "identity_links": [list(p) for p in links]}
        print(f"  {slug:<14} {len(keys):>4} entities  {len(links):>3} identity links  "
              f"{sum(1 for k in keys if gender.get(k)):>3} gendered  "
              f"{sum(1 for k in keys if not gender.get(k)):>4} ungendered")

    if ratio_rescued:
        print(f"\n  ratio-rejected but title-preceded, readmitted: {len(ratio_rescued)}")
        for k, (h, c, l) in sorted(ratio_rescued.items(), key=lambda kv: -kv[1][0]):
            print(f"    {k:<16} {h:>3} titled · {c:>4} cap / {l:>3} lower")

    if a.min_mid_ratio > 0:
        print(f"\n  mid-sentence filter (>= {a.min_mid} and >= {a.min_mid_ratio:.2f} of capitals): "
              f"dropped {len(mid_dropped)} candidates")
        for k, (m, t) in sorted(mid_dropped.items(), key=lambda kv: -kv[1][1])[:20]:
            print(f"    {k:<16} {m:>4} mid / {t:>4} caps")
        if len(mid_dropped) > 20:
            print(f"    ... and {len(mid_dropped) - 20} more")
        print(f"  rescued by honorific/possessive: {len(rescued)}")
        for k, (h, po) in sorted(rescued.items(), key=lambda kv: -sum(kv[1])):
            print(f"    {k:<16} {h:>3} titled · {po:>3} possessive")

    if neg_controls:
        print("\n  negative control -- these are NOT names and must be DROPPED:")
        for name in neg_controls:
            hits = [s for s, r in report.items() if name.lower() in r["entities"]]
            ok = not hits
            print(f"    [{'PASS' if ok else 'FAIL'}] {name:<14} "
                  f"{'dropped' if ok else 'STILL AN ENTITY in ' + ', '.join(hits)}")
            if not ok:
                failed_control.append(f"{name} (negative)")

    if controls:
        print("\n  positive control -- names known to be real must be FOUND:")
        for name in controls:
            hits = [s for s, r in report.items() if name.lower() in r["entities"]]
            ok = bool(hits)
            print(f"    [{'PASS' if ok else 'FAIL'}] {name:<14} {', '.join(hits) if hits else 'NOT DETECTED'}")
            if not ok:
                failed_control.append(name)

    if a.out:
        Path(a.out).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\n  wrote {a.out}")
    if failed_control:
        print(f"\n== POSITIVE CONTROL FAILED for {failed_control} -- the detector's negatives are worthless")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())