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

"""R49 Stage D2 (final) + D3 — entity resolution and deterministic rename augmentation.

D2's gender resolution is TITLE-FIRST, and that is the change from F02's method.
F02 used pronoun proximity and recorded that it is structurally blind to the
first-person narrator, whose name appears mainly in dialogue surrounded by other
people's pronouns. Measured here on Bronte, proximity called **Jane male** -- the
narrator of Jane Eyre, and the single worst entity to get wrong.

Titles do not have that blind spot. `Miss Eyre`, `Mrs. Fairfax`, `Mr. Rochester`,
`Madame Beck`, `M. Paul` are unambiguous and a 19th-century novel is saturated
with them. Measured: 16 entities resolved, **zero wrong**, with every ambiguous
case landing on HELD rather than on a guess -- shared family surnames like
Helstone and Pelet, which genuinely belong to both a man and a woman, hold as
they should.

Held is cheap; wrong is poison. **A HELD entity is simply not renamed.** An
un-renamed name costs a little augmentation; a mis-gendered one scrambles pronoun
agreement through every copy and nothing downstream would catch it.

Pool is French + English (operator, 2026-09-10), weighted per work by setting:
the Brussels novels draw more French, the Yorkshire novels more English. Locales
are restricted to fr_FR/fr_BE/en_GB/en_IE -- en_US and en_AU carry modern
surnames that are wrong register for the 1840s before any diacritic question.
"""
from __future__ import annotations
import argparse, collections, json, random, re, sys, unicodedata
from pathlib import Path

TOKEN = re.compile(r"[A-Za-zÀ-ÿŒœÆæ][A-Za-zà-ÿœæ\-]*")
MALE_T = r"(?:Mr|Sir|Master|Monsieur|M|Lord|Captain|Colonel|Major|Doctor|Dr|Reverend|King|Prince|Duke|Squire)"
FEM_T = r"(?:Mrs|Miss|Madame|Mme|Mademoiselle|Mlle|Lady|Madam|Queen|Princess|Duchess)"

FRENCH_LOCALES = ["fr_FR", "fr_BE"]
ENGLISH_LOCALES = ["en_GB", "en_IE"]
#: Brussels novels lean French, Yorkshire novels lean English. Register, not
#: orthography -- a Yorkshire mill town full of Parisian surnames reads wrong.
FRENCH_SHARE = {"villette": 0.60, "the-professor": 0.60, "jane-eyre": 0.25, "shirley": 0.25}

#: The pool is now per-corpus rather than per-author-hardcoded, because the same
#: register argument points somewhere else for every corpus. Brontë EXCLUDES en_US
#: (modern surnames read wrong for the 1840s); contemporary American romance wants
#: exactly those, with the European admixture F02 found matches Yarros's register.
#: Defaults reproduce the Brontë run byte-for-byte, so this is additive.
PRESETS = {
    "bronte": {"a": ("fr", FRENCH_LOCALES), "b": ("en", ENGLISH_LOCALES),
               "share": FRENCH_SHARE, "default_share": 0.25},
    "yarros": {"a": ("us", ["en_US", "en_CA"]),
               "b": ("eu", ["es_ES", "es_MX", "it_IT", "de_DE", "fr_FR"]),
               "share": {}, "default_share": 0.62},
    # Hemingway's cast is genuinely multilingual -- Spanish, Italian, Cuban and American
    # characters in the same books -- and the corpus alphabet agrees: 1,496 non-ASCII letters
    # across 23 forms, against Yarros's 2. The F02 rule (a pool's character inventory must be
    # a SUBSET of the corpus's) therefore permits accents here where it forbade them for
    # Yarros, and forbidding them would strand every Spanish and Italian name in the cast.
    # Share favours the romance-language pool because the two largest works by far -- For Whom
    # the Bell Tolls and Islands in the Stream -- are set in Spain and the Caribbean.
    "hemingway": {"a": ("us", ["en_US", "en_GB"]),
                  "b": ("rom", ["es_ES", "es_MX", "it_IT", "fr_FR"]),
                  "share": {}, "default_share": 0.45},
    # McCarthy is NOT a narrower Hemingway, and reusing that preset would have been the
    # easy wrong answer. Both are Spanish-inflected, but Hemingway's romance pool carries
    # it_IT and fr_FR for his Italian and French casts -- and McCarthy writes neither
    # language and has no such characters. Drawing from it would drop Italian and French
    # surnames into a Texas-Mexico border novel, which is exactly the register error the
    # per-corpus pool exists to prevent. en_GB goes for the same reason: Hemingway has
    # English characters (Brett Ashley, "the Englishman"), McCarthy's Anglo cast is
    # Texan and Tennessean throughout.
    # Share is even. The six works split about half and half: The Crossing and much of
    # All the Pretty Horses and Cities of the Plain are set in Mexico, while Blood
    # Meridian's gang and No Country's cast are Anglo. The corpus alphabet agrees that
    # accents belong -- 1,411 non-ASCII letters across 14 forms, all Spanish (á é í ñ ó ú ü).
    "mccarthy": {"a": ("us", ["en_US"]),
                 "b": ("mex", ["es_MX", "es_ES"]),
                 "share": {}, "default_share": 0.50},
}


def title_gender(text: str) -> dict[str, str]:
    mt = collections.Counter(m.group(1).lower() for m in
                             re.finditer(MALE_T + r"\.?\s+([A-ZÀ-Þ][a-zà-ÿœæ\-]+)", text))
    ft = collections.Counter(m.group(1).lower() for m in
                             re.finditer(FEM_T + r"\.?\s+([A-ZÀ-Þ][a-zà-ÿœæ\-]+)", text))
    out = {}
    for k in set(mt) | set(ft):
        M, F = mt[k], ft[k]
        if M >= 3 and M >= 3 * max(F, 1):
            out[k] = "m"
        elif F >= 3 and F >= 3 * max(M, 1):
            out[k] = "f"
    return out


def build_pool(dict_path: Path, alphabet: set[str], preset: str = "bronte") -> dict:
    d = json.loads(dict_path.read_text())
    pool = {}
    cfg = PRESETS[preset]
    for label, locales in (cfg["a"], cfg["b"]):
        m, f, s = set(), set(), set()
        for loc in locales:
            v = d["by_locale"].get(loc, {})
            m |= set(v.get("male", []))
            f |= set(v.get("female", []))
            for k in ("surnames_neutral", "surnames_male", "surnames_female"):
                s |= set(v.get(k, []))
        # ⚠ F02's subset rule, applied with Bronte's OWN alphabet rather than a
        # global ASCII fold: French accents are IN because she writes French
        # constantly; Czech/Latvian/Slovak marks are OUT because they never appear.
        keep = lambda n: n and n[:1].isupper() and all((not c.isalpha()) or c in alphabet for c in n)
        pool[label] = {"male": sorted(filter(keep, m)),
                       "female": sorted(filter(keep, f)),
                       "surname": sorted(filter(keep, s))}
    return pool


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("corpus")
    ap.add_argument("--entities", required=True)
    ap.add_argument("--dictionary", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--copies", type=int, default=6)
    ap.add_argument("--seed", type=int, default=4919)
    # ⚠ THE VAL SPLIT IS ONE CHAPTER INDEX PER WORK, so its SIZE is set by how many WORKS a
    # corpus has, not by how many words. That is why it takes a LIST now. Measured across the
    # line: Hemingway has 10 works -> 9 val units -> 36,563 words per copy and a decisive
    # gate; Brontë has 4 works -> 4 units -> 17,043 words and a gate that could not resolve
    # its own effect. A 588k-word corpus of 6 works would land at Brontë's end of that on a
    # single index, and no amount of corpus size fixes it.
    ap.add_argument("--holdout-chapter", type=int, nargs="+", default=[10],
                    help="chapter index/indices held out as val in EVERY work. Space them "
                         "apart -- adjacent chapters are more correlated with each other "
                         "than two drawn from different parts of a book.")
    ap.add_argument("--preset", default="bronte", choices=sorted(PRESETS),
                    help="which corpus's name-pool register to draw from")
    ap.add_argument("--scope", default="work", choices=("work", "corpus"),
                    help="`work` maps each work independently (reproduces the Bronte run); "
                         "`corpus` uses ONE map across every work in a copy")
    ap.add_argument("--phrase-map", default=None,
                    help="JSON with `phrases` (multiword) and `tokens` (capitalised single "
                         "words) neutralising in-world compounds the unigram pass cannot reach")
    ap.add_argument("--min-cap", type=int, default=8,
                    help="minimum capitalised count for an entity to be renamed; below it "
                         "the entity is left in the text verbatim")
    a = ap.parse_args()

    corpus = Path(a.corpus)
    man = json.loads((corpus / "manifest.json").read_text())
    alphabet = set(json.loads((corpus / "corpus_alphabet.json").read_text())["letters"])
    ents_all = json.loads(Path(a.entities).read_text())
    holdout = set(a.holdout_chapter)
    pool = build_pool(Path(a.dictionary), alphabet, a.preset)
    cfg = PRESETS[a.preset]
    label_a, label_b = cfg["a"][0], cfg["b"][0]
    # ⚠ Collision filter, against THIS corpus. F02 dropped 35 names for colliding
    # with the Yarros source so a rename could never map one of the author's
    # entities onto another; that filter is corpus-specific and does not carry.
    # Measured here before adding it: `Burns` and `Marie` were drawn as
    # replacements and are themselves Bronte entities, which reads as a leak in
    # the gate and is worse than it looks -- it silently merges two characters.
    source_names = {e["surface"] for w in ents_all.values() for e in w["entities"].values()}
    source_names |= {n.split()[0] for n in source_names if " " in n}
    dropped = 0
    for lang in pool:
        for bucket in pool[lang]:
            before = len(pool[lang][bucket])
            # ⚠ By COMPONENT, not by whole string. Measured: the pool drew the
            # compound `Pierre-Yves` while `Pierre` (Mademoiselle St. Pierre) is a
            # Villette character, so a whole-string comparison passed it and the
            # leak gate then matched the component. The original was correctly
            # renamed -- it is not a leak -- but a replacement sharing a component
            # with a source character invites exactly the conflation the rename
            # exists to prevent.
            pool[lang][bucket] = [
                n for n in pool[lang][bucket]
                if n not in source_names
                and not (set(re.split(r"[-\s’']", n)) & source_names)]
            dropped += before - len(pool[lang][bucket])
    print(f"  collision filter: dropped {dropped} pool names that collide with "
          f"{len(source_names)} source entities in THIS corpus ({a.preset})")
    print("  pool (alphabet-filtered): " + "  ".join(
        f"{lab} {len(pool[lab]['male'])}m/{len(pool[lab]['female'])}f/{len(pool[lab]['surname'])}s"
        for lab in (label_a, label_b)))

    # ⚠ Applied AFTER the entity substitution, so it can never eat a replacement
    # name. Multiword first and longest first; single tokens are case-SENSITIVE
    # and whole-word, so a dragon's lowercase `wing` survives while `Fourth Wing`
    # does not.
    phrase_sub = None
    if a.phrase_map:
        pm = json.loads(Path(a.phrase_map).read_text())
        table = {**pm.get("phrases", {}), **pm.get("tokens", {})}
        if table:
            pat_p = re.compile(r"\b(" + "|".join(re.escape(k) for k in
                                                 sorted(table, key=len, reverse=True)) + r")\b")
            phrase_sub = lambda t: pat_p.sub(lambda m: table[m.group(1)], t)
            print(f"  phrase map {a.phrase_map}: {len(pm.get('phrases', {}))} phrases + "
                  f"{len(pm.get('tokens', {}))} capitalised tokens")

    out = Path(a.out); (out / "copies").mkdir(parents=True, exist_ok=True)
    stats = {"copies": a.copies, "seed": a.seed, "works": {}, "renamed": 0, "held": 0}

    works = {}
    for w in man["works"]:
        rows = [json.loads(l) for l in (corpus / w["path"]).read_text(encoding="utf-8").splitlines()]
        works[w["slug"]] = rows

    # ---- D2 final: decide, per work, which entities are renameable ----------
    plans = {}
    for slug, rows in works.items():
        text = "\n\n".join(r["text"] for r in rows)
        tg = title_gender(text)
        ents = ents_all[slug]["entities"]
        titled = set(tg)
        renameable, held = {}, []
        for key, e in ents.items():
            if "’" in key or "'" in key or e["cap"] < a.min_cap:
                continue                      # possessives/contractions are not entities
            g = tg.get(key) or e.get("gender")
            if g:
                renameable[key] = {"surface": e["surface"], "kind": "given", "gender": g}
            else:
                # ⚠ Everything else is STILL renamed -- from the gender-NEUTRAL
                # surname/place pool. The operator's Yarros directive was "rename
                # all proper nouns", and holding a place leaks it: `Thornfield`
                # appears 100 times in Jane Eyre and is as author-specific as
                # `Riders Quadrant` was. Substituting a neutral token makes NO
                # gender claim, so no gender claim can be wrong -- the prose keeps
                # whatever pronoun it already had. Held-means-ungendered, not
                # held-means-unrenamed.
                renameable[key] = {"surface": e["surface"], "kind": "surname", "gender": None}
                held.append(key)
        plans[slug] = renameable
        stats["works"][slug] = {"renamed": len(renameable), "gendered": len(renameable)-len(held),
                                "neutral": len(held)}
        stats["renamed"] += len(renameable); stats["held"] += len(held)
        print(f"  {slug:<14} renamed {len(renameable):>3}  ({len(renameable)-len(held)} gendered, {len(held)} neutral)")

    # ⚠ CORPUS SCOPE. Per-work maps leak across works and this is measurable, not
    # theoretical: `Rebel` is detected in `rebel` and renamed there, then printed
    # verbatim in the two Renegades books where it sits below threshold. A
    # whole-corpus gate catches it; a per-work one reports clean. It also fixes a
    # thing the Bronte corpus never had to care about -- Yarros is TWO SERIES, so
    # Violet has to be the same person in Fourth Wing and Iron Flame, and a
    # per-work draw gives her two different names inside one copy.
    if a.scope == "corpus":
        merged: dict[str, dict] = {}
        genders: dict[str, set] = collections.defaultdict(set)
        for slug, plan in plans.items():
            for key, v in plan.items():
                merged.setdefault(key, {"surface": v["surface"], "kind": v["kind"], "gender": None})
                if v["gender"]:
                    genders[key].add(v["gender"])
        conflicts = 0
        for key, v in merged.items():
            g = genders.get(key, set())
            if len(g) == 1:
                v["gender"] = next(iter(g)); v["kind"] = "given"
            else:
                if len(g) > 1:
                    conflicts += 1
                v["kind"] = "surname"         # held -> neutral pool, still renamed
        n_gendered = sum(1 for v in merged.values() if v["gender"])
        print(f"  corpus scope: {len(merged)} distinct surfaces "
              f"({n_gendered} gendered, {len(merged) - n_gendered} neutral), "
              f"{conflicts} gender conflicts held")
        plans = {slug: merged for slug in plans}

    # ---- D3: N seeded copies, one consistent map per copy -------------------
    emitted = 0
    for c in range(a.copies):
        rng = random.Random(a.seed + c * 1000)
        corpus_map, corpus_used = {}, set()
        for slug, rows in works.items():
            # Share of pool A for this work. Brontë sets it per novel (Brussels
            # vs Yorkshire); Yarros uses one default, because the register does
            # not split by book the way hers does.
            share_a = cfg["share"].get(slug, cfg["default_share"])
            used = corpus_used if a.scope == "corpus" else set()

            def draw(kind: str, gender: str | None) -> str:
                lang = label_a if rng.random() < share_a else label_b
                bucket = {"m": "male", "f": "female"}.get(gender or "", "surname")
                for _ in range(200):
                    n = rng.choice(pool[lang][bucket])
                    if n not in used:
                        used.add(n); return n
                return rng.choice(pool[lang][bucket])

            if a.scope == "corpus":
                for k, v in plans[slug].items():
                    corpus_map.setdefault(k, draw(v["kind"], v["gender"]))
                mapping = corpus_map
            else:
                mapping = {k: draw(v["kind"], v["gender"]) for k, v in plans[slug].items()}
            pat = re.compile(r"\b(" + "|".join(sorted((re.escape(v["surface"]) for v in plans[slug].values()),
                                                      key=len, reverse=True)) + r")\b")
            surf2key = {v["surface"]: k for k, v in plans[slug].items()}
            path = out / "copies" / f"{slug}.copy{c}.jsonl"
            with path.open("w", encoding="utf-8") as fh:
                for r in rows:
                    txt = pat.sub(lambda m: mapping[surf2key[m.group(1)]], r["text"])
                    if phrase_sub:
                        txt = phrase_sub(txt)
                    split = "val" if r["chapter"] in holdout else "train"
                    fh.write(json.dumps({"work": slug, "copy": c, "chapter": r["chapter"],
                                         "split": split, "text": txt}, ensure_ascii=False) + "\n")
                    emitted += 1
        print(f"  copy {c}: written")

    (out / "rename_stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2))
    print(f"\n  {emitted:,} chapter-records across {a.copies} copies -> {out}")
    return 0


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