phase1/scripts/backfill_headnotes.py

#!/usr/bin/env python3
"""Backfill synthetic headnotes for the 1970s-80s crater (~5,900 docs).

Reporter headnote coverage collapses to 1.6% (1970s) / 3.8% (1980s) — exactly the
golden era of constitutional doctrine. Downstream, the skim gate and judge ground on
`held`-or-first-pages, so crater docs present a cover page and silently exit the
pipeline. This writes an LLM-drafted issue/held in reporter register from each
judgment's own text, marked synthetic:true (the UI must disclose it).

Selection: held_len < 200 AND decision_year in [1965, 1995] (from corpus_ledger.jsonl).
Model: deepseek-v4-flash (non-thinking pinned), temperature 0. ~$6-7 for the full set. Requires DEEPSEEK_API_KEY in phase1/scripts/.env
or the environment.

Run:   python phase1/scripts/backfill_headnotes.py [data_dir] [--dry-run N] [--limit N]
Out:   <data_dir>/synthetic_headnotes.jsonl  {doc_id, issue, held, synthetic: true}
       (appends; already-done doc_ids are skipped -> resumable)
After: re-run build_held_vectors.py to fold synthetic helds into the HELD arm.
"""
import json, os, re, sys, time

data_dir = next((a for a in sys.argv[1:] if not a.startswith("--")), None) \
    or os.environ.get("THEMIS_DATA", "phase1/data/thor_artifacts")
DRY = 0
if "--dry-run" in sys.argv:
    i = sys.argv.index("--dry-run"); DRY = int(sys.argv[i + 1]) if len(sys.argv) > i + 1 else 3
LIMIT = int(sys.argv[sys.argv.index("--limit") + 1]) if "--limit" in sys.argv else None
MODEL = os.environ.get("THEMIS_LLM_MODEL", "deepseek-v4-flash")   # deepseek-chat alias dies 2026-07-24

SYS = ("You are a Supreme Court of India law reporter writing an eSCR-style headnote from the "
       "judgment text supplied. Output STRICT JSON: {\"issue\": \"...\", \"held\": \"...\"}. "
       "'issue' = the question(s) of law before the Court, 1-3 sentences. 'held' = what the Court "
       "decided and its reasoning, 150-400 words, neutral reporter register, past tense "
       "(\"Held: ...\"), naming doctrines and provisions precisely. Use ONLY the supplied text; "
       "if it is insufficient, output {\"issue\": \"\", \"held\": \"\"}.")

def select_docs():
    ledger = {}
    for line in open(os.path.join(data_dir, "corpus_ledger.jsonl"), encoding="utf-8"):
        r = json.loads(line); ledger[r["doc_id"]] = r
    picks = [d for d, r in ledger.items()
             if (r.get("held_len") or 0) < 200
             and r.get("decision_year") and 1965 <= r["decision_year"] <= 1995
             and r.get("text_health") == "ok"]
    return picks

def doc_text(doc_ids):
    """Front ~10 pages; for LONG judgments also the tail — the operative holding of a
    200-page judgment (e.g. Bachan Singh's 'rarest of rare' at para 209) lives at the
    END and a front-only summary misses the doctrine the case is famous for."""
    want = set(doc_ids); front = {d: [] for d in doc_ids}; tail = {d: [] for d in doc_ids}
    nch = {d: 0 for d in doc_ids}
    pat = re.compile(r'"doc_id":\s*"([^"]+)"')
    for line in open(os.path.join(data_dir, "escr_chunks.jsonl"), encoding="utf-8"):
        d = pat.search(line[:120]).group(1)
        if d in want:
            t = json.loads(line)["text"]; nch[d] += 1
            if len(front[d]) < 10: front[d].append(t)
            tail[d].append(t)
            if len(tail[d]) > 5: tail[d].pop(0)              # rolling last-5 window
    out = {}
    for d in doc_ids:
        if nch[d] > 50:                                       # long doc: head + tail pack
            out[d] = (" ".join(front[d])[:10000] + "\n[... middle omitted ...]\n"
                      + " ".join(tail[d])[:5500])
        else:
            out[d] = " ".join(front[d] + [c for c in tail[d] if c not in front[d]])[:16000]
    return out

def main():
    picks = select_docs()
    outp = os.path.join(data_dir, "synthetic_headnotes.jsonl")
    done = set()
    if os.path.exists(outp):
        for line in open(outp, encoding="utf-8"):
            done.add(json.loads(line)["doc_id"])
    todo = [d for d in picks if d not in done]
    if LIMIT: todo = todo[:LIMIT]
    print(f"[backfill] crater docs: {len(picks)} | done: {len(done)} | todo: {len(todo)}", flush=True)

    if DRY:
        texts = doc_text(todo[:DRY])
        for d in todo[:DRY]:
            print(f"\n--- DRY {d} --- prompt head:\n{texts.get(d,'')[:500]}", flush=True)
        print(f"\n[backfill] dry-run only ({DRY} docs shown); no API calls made.", flush=True)
        return

    key = os.environ.get("DEEPSEEK_API_KEY", "")
    if not key:
        env = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
        if os.path.exists(env):
            for l in open(env):
                if l.startswith("DEEPSEEK_API_KEY="): key = l.split("=", 1)[1].strip()
    if not key:
        sys.exit("[backfill] DEEPSEEK_API_KEY missing (env or phase1/scripts/.env) — aborting.")
    import threading
    from concurrent.futures import ThreadPoolExecutor

    import requests
    WORKERS = int(os.environ.get("THEMIS_BACKFILL_WORKERS", "8"))
    lock = threading.Lock()
    done_n = [0, 0]                                       # ok, skip

    def one(d, t):
        try:
            r = requests.post("https://api.deepseek.com/chat/completions",
                headers={"Authorization": f"Bearer {key}"},
                json={"model": MODEL, "temperature": 0,
                      "thinking": {"type": "disabled"},
                      "response_format": {"type": "json_object"},
                      "messages": [{"role": "system", "content": SYS},
                                   {"role": "user", "content": t}]},
                timeout=180)
            j = json.loads(r.json()["choices"][0]["message"]["content"])
            if len(j.get("held") or "") > 100:
                row = json.dumps({"doc_id": d, "issue": j.get("issue", ""),
                                  "held": j["held"], "synthetic": True, "model": MODEL},
                                 ensure_ascii=False)
                with lock:
                    f.write(row + "\n"); f.flush(); done_n[0] += 1
                return
        except Exception as e:
            print(f"[backfill] {d}: {e}", flush=True); time.sleep(2)
        with lock:
            done_n[1] += 1

    B = 200
    t0 = time.time()
    with open(outp, "a", encoding="utf-8") as f:
        for s in range(0, len(todo), B):
            batch = todo[s:s + B]
            texts = doc_text(batch)                       # one chunks-file pass per 200 docs
            jobs = [(d, texts[d]) for d in batch if len(texts.get(d, "")) >= 1500]
            with ThreadPoolExecutor(WORKERS) as ex:
                list(ex.map(lambda a: one(*a), jobs))
            el = time.time() - t0
            print(f"[backfill] {min(s+B,len(todo))}/{len(todo)} ok={done_n[0]} skip={done_n[1]} "
                  f"({done_n[0]/el*3600:.0f}/h)", flush=True)

if __name__ == "__main__":
    main()