extract_check.py

#!/usr/bin/env python3
"""Extract-then-check: the small-model path for the hardest error class.

A small model is bad at multi-step arithmetic across scattered facts, and good at
pulling facts out of text. So we ask it only to EXTRACT dated/aged facts (with quotes),
then Python computes the contradictions. The reasoning lives in code, where it is
deterministic, instant, free, and never hallucinated.

This targets the age/timeline class: the exact bug an LLM-reasoning pass just missed.

  python extract_check.py MANUSCRIPT.txt [--model gemma3:27b]
"""
from __future__ import annotations
import argparse, json, re, sys
from throughline.semantic import Model, extract_json, Grounder

# ------------------------------------------------------- number words -> ints
_ONES = {"zero":0,"one":1,"two":2,"three":3,"four":4,"five":5,"six":6,"seven":7,"eight":8,
         "nine":9,"ten":10,"eleven":11,"twelve":12,"thirteen":13,"fourteen":14,"fifteen":15,
         "sixteen":16,"seventeen":17,"eighteen":18,"nineteen":19}
_TENS = {"twenty":20,"thirty":30,"forty":40,"fifty":50,"sixty":60,"seventy":70,"eighty":80,"ninety":90}

def word_to_int(s):
    if s is None: return None
    s = str(s).strip().lower().replace("–","-")
    if re.fullmatch(r"-?\d+", s): return int(s)
    s = s.replace("-", " ")
    toks = s.split()
    total = 0; ok = False
    for t in toks:
        if t in _ONES: total += _ONES[t]; ok = True
        elif t in _TENS: total += _TENS[t]; ok = True
        elif t == "hundred": total = (total or 1) * 100; ok = True
        elif t == "thousand": total = (total or 1) * 1000; ok = True
    return total if ok else None

def parse_year(s):
    """'oh-nine'->2009, '09'->2009, 'two thousand nine'->2009, 2009->2009, 1994->1994."""
    if s is None: return None
    s = str(s).strip().lower().replace("’","'")
    m = re.search(r"\b(19|20)\d{2}\b", s)
    if m: return int(m.group(0))
    m = re.search(r"(?:oh|aught|nought|')[ -]?(\w+)", s)   # "oh-nine", "'09"
    if m:
        n = word_to_int(m.group(1)) if not m.group(1).isdigit() else int(m.group(1))
        if n is not None and 0 <= n <= 30: return 2000 + n
    n = word_to_int(s)
    if n is not None and 0 <= n <= 30: return 2000 + n
    if n is not None and 30 < n <= 99: return 1900 + n
    return None

# ------------------------------------------------------- extraction
SCHEMA = """Output a JSON array of temporal facts. Each fact fixes a person's age or an event's date.
Fields: {"person": name or null, "subject": event/thing the number is about or null,
 "kind": one of ["age_now","age_at_year","made_at_age","event_year","span_years"],
 "number": the number exactly as written, "year": the year phrase if the fact names one else null,
 "quote": the exact sentence}
Guide:
 age_now      "Rhett is thirty"                      -> person=Rhett, number=thirty
 age_at_year  "I was twenty-two in oh-nine"          -> person, number=twenty-two, year=oh-nine
 made_at_age  "wrote the first song at nineteen"     -> person, subject=song, number=nineteen
 event_year   "the wedding was oh-nine"              -> subject=wedding, year=oh-nine
 span_years   "held eleven years" / "fifteen years back" -> subject=(the thing), number=eleven
Only facts explicitly in the text. Do not infer. Return ONLY the JSON array."""

_NUMWORD = r"(?:twenty-\w+|thirty-\w+|forty-\w+|fifty-\w+|sixty-\w+|" \
           r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|eighteen|nineteen|" \
           r"one|two|three|four|five|six|seven|eight|nine|ten|" \
           r"twenty|thirty|forty|fifty|sixty|seventy|eighty|ninety|\d{1,3})"
_YEARWORD = r"(?:oh[- ]\w+|aught[- ]\w+|['’]\d{2}|\b(?:19|20)\d{2}\b)"

def regex_atoms(text):
    """Deterministic backbone for the killer age/date patterns, both word orders, with the
    two spans that matter tagged by their context (the made-thing's 'held N years', and the
    tally 'N years' since an event). The code does the math; these just feed it clean atoms."""
    out = []
    # age_at_year, BOTH orders: "was 22 ... oh-nine"  AND  "oh-nine ... I was 22"
    for m in re.finditer(rf"\bwas\s+({_NUMWORD})\b[^.!\n]{{0,25}}?\b({_YEARWORD})", text, re.I):
        out.append({"person": None, "kind": "age_at_year", "number": m.group(1), "year": m.group(2), "quote": m.group(0)})
    for m in re.finditer(rf"\b({_YEARWORD})\b[^.!\n]{{0,25}}?\bwas\s+({_NUMWORD})\b", text, re.I):
        out.append({"person": None, "kind": "age_at_year", "number": m.group(2), "year": m.group(1), "quote": m.group(0)})
    # made a thing "at NINETEEN"
    for m in re.finditer(rf"\b(?:wrote|first (?:one|song)|the first)\b[^.?!]{{0,45}}?\bat\s+({_NUMWORD})\b", text, re.I):
        out.append({"person": None, "kind": "made_at_age", "subject": "song", "number": m.group(1), "year": None, "quote": m.group(0)})
    # the made-thing's own duration: "held eleven years", "eleven years of"
    for m in re.finditer(rf"\b(?:held|kept|it['’]s held|for)\s+({_NUMWORD})\s+years?\b", text, re.I):
        out.append({"person": None, "kind": "held_span", "subject": "song", "number": m.group(1), "year": None, "quote": m.group(0)})
    # a tally measured from an event: "fifteen years", near since/working/conspiracy
    for m in re.finditer(rf"\b({_NUMWORD})\s+years\b[^.?!]{{0,30}}?(?:working|figured|since|conspiracy|tally)"
                         rf"|(?:since|conspiracy|tally|working)[^.?!]{{0,30}}?\b({_NUMWORD})\s+years\b", text, re.I):
        out.append({"person": None, "kind": "tally_span", "number": m.group(1) or m.group(2), "year": None, "quote": m.group(0)})
    return out

def extract_atoms(mdl, manuscript):
    sys_p = ("You extract temporal facts from prose into JSON. An AGE is how old a PERSON is in years; "
             "a COUNT of objects (e.g. 'forty-some songs') is NOT an age, never record it as one. Always "
             "name the 'person' the age belongs to (in dialogue, the speaker; in narration, the point-of-view "
             "character). Be literal; do not infer years that are not written.")
    out = mdl.ask(sys_p, SCHEMA + "\n\nTEXT:\n" + manuscript, max_tokens=1800)
    atoms = [a for a in extract_json(out, "array")
             if isinstance(a, dict) and a.get("kind") and a.get("quote")]
    # union with the deterministic regex backbone, deduped by (kind, number, year)
    seen = {(a.get("kind"), str(a.get("number")), str(a.get("year"))) for a in atoms}
    for r in regex_atoms(manuscript):
        k = (r["kind"], str(r["number"]), str(r["year"]))
        if k not in seen:
            seen.add(k); atoms.append(r)
    return atoms

# ------------------------------------------------------- deterministic check
def _nums(atoms, kind):
    return [(word_to_int(a.get("number")), a) for a in atoms if a["kind"] == kind and word_to_int(a.get("number"))]

def resolve_present(atoms):
    """present = the year an age_at_year/event names, plus the tally span measured from it."""
    years = [parse_year(a.get("year")) for a in atoms if a["kind"] in ("event_year", "age_at_year")]
    years = [y for y in years if y]
    tallies = [n for n, _ in _nums(atoms, "tally_span")] or [n for n, _ in _nums(atoms, "span_years")]
    if years and tallies:
        return max(years) + max(tallies)
    return None

def derive_ages(atoms, present):
    """Independent current-age derivations, bucketed by person ('(protagonist)' when unnamed)."""
    out = {}
    def add(p, age, q, path):
        if age is not None:
            out.setdefault((p or "(protagonist)").strip().lower(), []).append((age, q, path))
    held = [n for n, _ in _nums(atoms, "held_span")]
    for a in atoms:
        n = word_to_int(a.get("number"))
        if a["kind"] == "age_now":
            add(a.get("person"), n, a["quote"], "stated age")
        elif a["kind"] == "age_at_year" and present:
            y = parse_year(a.get("year"))
            if y and n is not None:
                add(a.get("person"), n + (present - y), a["quote"], f"age {n} in {y}, present {present}")
        elif a["kind"] == "made_at_age" and n is not None and held:
            add(a.get("person"), n + min(held), a["quote"], f"made at {n} + held {min(held)} years")
    return out

def check(atoms, tol=2):
    present = resolve_present(atoms)
    ages = derive_ages(atoms, present)
    findings = []
    for person, ests in ages.items():
        if len(ests) < 2: continue
        lo = min(ests, key=lambda e: e[0]); hi = max(ests, key=lambda e: e[0])
        if hi[0] - lo[0] > tol:
            findings.append({
                "person": person, "spread": hi[0] - lo[0],
                "age_a": lo[0], "why_a": lo[2], "quote_a": lo[1],
                "age_b": hi[0], "why_b": hi[2], "quote_b": hi[1],
            })
    return present, ages, findings

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("manuscript")
    ap.add_argument("--model", default="gemma3:27b")
    ap.add_argument("--show-atoms", action="store_true")
    a = ap.parse_args()
    ms = open(a.manuscript, encoding="utf-8", errors="ignore").read()
    mdl = Model("ollama", a.model)
    ground = Grounder(ms)

    print(f"extract-then-check  [{a.model}]  {a.manuscript}")
    atoms = extract_atoms(mdl, ms)
    grounded = [x for x in atoms if ground.has(x["quote"])]
    print(f"  extracted {len(atoms)} atoms, {len(grounded)} grounded")
    if a.show_atoms:
        for x in grounded:
            print(f"    {x['kind']:12} person={x.get('person')} subj={x.get('subject')} "
                  f"num={x.get('number')} year={x.get('year')}")

    present, ages, findings = check(grounded)
    print(f"  present year resolved to: {present}")
    for person, ests in ages.items():
        vals = ", ".join(f"{a_}({p})" for a_, _, p in ests)
        print(f"  {person}: current-age derivations -> {vals}")
    print("=" * 66)
    if not findings:
        print("NO age/timeline contradiction found.")
    for f in findings:
        print(f"\nCONTRADICTION — {f['person']}: derived age {f['age_a']} vs {f['age_b']} "
              f"(off by {f['spread']})")
        print(f"  [{f['age_a']}] {f['why_a']}\n     \"{f['quote_a'][:90]}\"")
        print(f"  [{f['age_b']}] {f['why_b']}\n     \"{f['quote_b'][:90]}\"")

if __name__ == "__main__":
    main()