SPB Git

spb/firc0001_uqo

Public
12.6 KB · 149 lines python
1"""N11 - sampling frames and design for the human annotation pilot (cluster A, ECON-04). No labels produced.23Strata: item group x run x truncation bin (N5 shown fraction) x productive x confidence bin x category group;4productive sections n = 385 allocated proportionally with a floor of 10 and 2x oversampling of the 0.45-0.7 bin;5zero-mention sections drawn with inclusion probabilities proportional to (gate pass, log length) with6inverse-probability weights; seeded (20261105); development / holdout split. Frames with full text and7_focus(text) stay on the node (~/firc0001_review/frames/); the worktree copy truncates text to 2,000 chars and8carries no model output (the model's mentions are sealed in a separate file).9"""10from __future__ import annotations11import json, sys, math12from pathlib import Path13import numpy as np14sys.path.insert(0, str(Path(__file__).resolve().parent))15from common import *16from spid.extract.llm import _focus1718N_PROD = 385; N_ZERO = 385; FLOOR = 1019rng = np.random.default_rng(SEED)20r = Run("annotation_frame")21FR = WORK / "frames"; FR.mkdir(exist_ok=True)22con = connect(SPID_DB)23tmin = con.execute("select min(extracted_at) from extracted_sections").fetchone()[0]24feat = RESULTS / "focus" / "section_features.parquet"25con.execute(f"CREATE TEMP TABLE ft AS SELECT section_id, nchar, gate_py, shown_chars*1.0/nchar shown_frac FROM read_parquet('{feat}')")26con.execute(f"""CREATE TEMP TABLE S AS27SELECT e.section_id, e.accession_number, e.n_mentions, CASE WHEN e.extracted_at = TIMESTAMP '{tmin}' THEN 'june' ELSE 'october' END run,28       f.form_type, f.cik, f.ticker, f.company_name, f.filing_date, m.section_name item, m.word_count, ft.nchar, ft.gate_py gate, ft.shown_frac,29       CASE WHEN f.form_type='10-K' AND m.section_name='md_and_a' THEN '10-K md_and_a' WHEN f.form_type='10-K' THEN '10-K other'30            WHEN f.form_type='10-Q' AND m.section_name='md_and_a' THEN '10-Q md_and_a' WHEN f.form_type='10-Q' THEN '10-Q other' ELSE '8-K' END item_group,31       CASE WHEN ft.shown_frac >= 0.9 THEN 'shown>=0.9' WHEN ft.shown_frac >= 0.25 THEN 'shown0.25-0.9' ELSE 'shown<0.25' END trunc_bin,32       pm.max_conf, CASE WHEN pm.max_conf IS NULL THEN 'none' WHEN pm.max_conf <= 0.7 THEN 'conf0.45-0.7' ELSE 'conf>0.7' END conf_bin,33       CASE WHEN pm.types IS NULL THEN 'none' WHEN list_contains(pm.types, 'ma_integration') THEN 'ma_integration'34            WHEN len(list_intersect(pm.types, ['capex_program','plant_construction','manufacturing_expansion','infrastructure','data_center'])) > 0 THEN 'capex-type'35            WHEN len(list_intersect(pm.types, ['technology_deployment','ai_initiative','digital_transformation','erp_implementation','cloud_migration','automation'])) > 0 THEN 'technology-type'36            ELSE 'other' END cat_group37FROM extracted_sections e JOIN src_filings f ON f.accession_number = e.accession_number38JOIN (SELECT section_id, section_name, word_count FROM src_sections UNION ALL SELECT section_id, section_name, word_count FROM spid_sections) m USING (section_id)39LEFT JOIN ft USING (section_id)40LEFT JOIN (SELECT section_id, max(confidence) max_conf, list_distinct(list(project_type)) types FROM project_mentions GROUP BY 1) pm USING (section_id)""")41rules = RESULTS / "rules" / "rules_mentions.parquet"42if rules.exists():43    con.execute(f"CREATE TEMP TABLE rl AS SELECT section_id, rules_n FROM read_parquet('{rules}')")44    con.execute("ALTER TABLE S ADD COLUMN rules_n INTEGER"); con.execute("UPDATE S SET rules_n = rl.rules_n FROM rl WHERE S.section_id = rl.section_id")45else:46    con.execute("ALTER TABLE S ADD COLUMN rules_n INTEGER")47# ---- productive strata and allocation ---------------------------------------------------------------------------48strata = con.execute("""SELECT item_group, run, trunc_bin, conf_bin, cat_group, count(*) N FROM S WHERE n_mentions > 0 GROUP BY 1,2,3,4,5 ORDER BY 1,2,3,4,5""").fetchall()49Ntot = sum(s[5] for s in strata)50alloc = []51for ig, rn, tb, cb, cg, N in strata:52    share = N / Ntot * (2.0 if cb == "conf0.45-0.7" else 1.0)53    n = max(FLOOR, round(N_PROD * share)); n = min(n, N)54    alloc.append((ig, rn, tb, cb, cg, N, n, N / n))55n_prod_actual = sum(a[6] for a in alloc)56r.write_csv("design_productive_strata.csv", ["item_group", "run", "trunc_bin", "conf_bin", "cat_group", "N_population", "n_sampled", "weight"], alloc)57# ---- draw productive --------------------------------------------------------------------------------------------58rows = []59for ig, rn, tb, cb, cg, N, n, w in alloc:60    ids = con.execute("SELECT section_id FROM S WHERE n_mentions > 0 AND item_group=? AND run=? AND trunc_bin=? AND conf_bin=? AND cat_group=? ORDER BY section_id", [ig, rn, tb, cb, cg]).fetchall()61    pick = rng.choice(len(ids), size=n, replace=False)62    for i in pick:63        rows.append((ids[i][0], "productive", f"{ig}|{rn}|{tb}|{cb}|{cg}", w))64# ---- draw zero-mention with inclusion probabilities ∝ (gate, log length) -------------------------------------------65zero = con.execute("SELECT section_id, gate, nchar, item_group, run, trunc_bin FROM S WHERE n_mentions = 0 ORDER BY section_id").fetchall()66score = np.array([(2.0 if g else 1.0) * (1.0 + math.log10(max(30, nc))) for _, g, nc, *_ in zero])67p = score / score.sum()68pick = rng.choice(len(zero), size=N_ZERO, replace=False, p=p)69pi = N_ZERO * p   # first-order inclusion probability (approximation, sampling fraction 4e-4)70for i in pick:71    sid, g, nc, ig, rn, tb = zero[i]72    rows.append((sid, "zero", f"{ig}|{rn}|{tb}|gate={'pass' if g else 'reject'}", 1.0 / pi[i]))73con.execute("CREATE TEMP TABLE pick AS SELECT * FROM (VALUES " + ",".join(["(?,?,?,?)"] * len(rows)) + ") t(section_id, frame, stratum, weight)", [x for row in rows for x in row])74# dev / holdout split, stratified by frame x stratum, seeded75sel = con.execute("SELECT p.section_id, p.frame, p.stratum, p.weight, s.accession_number, s.form_type, s.item, s.cik, s.ticker, s.company_name, s.filing_date, s.run, s.word_count, s.nchar, s.gate, s.shown_frac, s.n_mentions, s.max_conf, s.cat_group, s.rules_n FROM pick p JOIN S s USING (section_id) ORDER BY p.frame, p.stratum, p.section_id").fetchall()76split = {}77by = {}78for x in sel: by.setdefault((x[1], x[2]), []).append(x[0])79for k, ids in by.items():80    perm = rng.permutation(len(ids))81    for j, i in enumerate(perm):82        split[ids[i]] = "dev" if j % 2 == 0 else "holdout"83cols = ["section_id", "frame", "stratum", "weight", "accession_number", "form_type", "item", "cik", "ticker", "company_name", "filing_date", "run", "word_count", "nchar", "gate_pass", "shown_frac", "n_mentions_model", "max_conf_model", "cat_group_model", "rules_n"]84# nested core subsample (first annotation pass): ~385 productive sections, proportional to the stratum population, min 1, seeded85core = set()86prod_by = {}87for x in sel:88    if x[1] == "productive": prod_by.setdefault(x[2], []).append(x[0])89for st, ids in prod_by.items():90    N_h = next(a[5] for a in alloc if f"{a[0]}|{a[1]}|{a[2]}|{a[3]}|{a[4]}" == st)91    k = max(1, min(len(ids), round(385 * N_h / Ntot)))92    for i in rng.choice(len(ids), size=k, replace=False): core.add(ids[i])93for x in sel:94    if x[1] == "zero": core.add(x[0])95r.write_csv("weights.csv", ["section_id", "frame", "stratum", "weight", "split", "core"], [(x[0], x[1], x[2], round(x[3], 4), split[x[0]], x[0] in core) for x in sel])96# worktree copies: no model output, text truncated to 2,000 chars97con.execute("CREATE TEMP TABLE txt AS SELECT t.section_id, t.text FROM (SELECT section_id, text FROM src_sections UNION ALL SELECT section_id, text FROM spid_sections) t JOIN pick USING (section_id)")98texts = dict(con.execute("SELECT section_id, text FROM txt").fetchall())99pub_cols = ["section_id", "split", "core", "frame", "stratum", "weight", "accession_number", "form_type", "item", "cik", "ticker", "company_name", "filing_date", "run", "word_count", "nchar", "gate_pass", "shown_frac", "text_2000"]100pub = [(x[0], split[x[0]], x[0] in core, x[1], x[2], round(x[3], 4), x[4], x[5], x[6], x[7], x[8], x[9], str(x[10]), x[11], x[12], x[13], x[14], round(x[15], 4) if x[15] is not None else None, (texts.get(x[0]) or "")[:2000]) for x in sel]101r.write_csv("frame_dev.csv", pub_cols, [p_ for p_ in pub if p_[1] == "dev"])102r.write_csv("frame_holdout.csv", pub_cols, [p_ for p_ in pub if p_[1] == "holdout"])103# node-only full frames + sealed model output104men = {}105for row in con.execute("SELECT section_id, mention_id, project_name, project_type, description, status, amount_usd, locations, partners, years, confidence FROM project_mentions WHERE section_id IN (SELECT section_id FROM pick)").fetchall():106    men.setdefault(row[0], []).append(dict(zip(["mention_id", "project_name", "project_type", "description", "status", "amount_usd", "locations", "partners", "years", "confidence"], row[1:])))107for name in ("dev", "holdout"):108    with open(FR / f"frame_{name}.jsonl", "w") as fh:109        for x in sel:110            if split[x[0]] != name: continue111            t = texts.get(x[0]) or ""112            fh.write(json.dumps({"section_id": x[0], "frame": x[1], "stratum": x[2], "weight": x[3], "accession_number": x[4], "form_type": x[5], "item": x[6], "cik": x[7], "ticker": x[8],113                                 "company_name": x[9], "filing_date": str(x[10]), "run": x[11], "word_count": x[12], "nchar": x[13], "gate_pass": x[14], "shown_frac": x[15],114                                 "text": t, "focus_text_seen_by_model": ("SECTION: " + str(x[6]) + "\n\n" + _focus(t))[:14000]}, default=str) + "\n")115with open(FR / "model_output_SEALED.jsonl", "w") as fh:116    for x in sel:117        fh.write(json.dumps({"section_id": x[0], "split": split[x[0]], "n_mentions": x[16], "mentions": men.get(x[0], [])}, default=str) + "\n")118import shutil119shutil.copy(r.out / "weights.csv", FR / "weights.csv")120# ---- design numbers: CI widths and MDE ---------------------------------------------------------------------------------121w_prod = np.array([x[3] for x in sel if x[1] == "productive"]); w_zero = np.array([x[3] for x in sel if x[1] == "zero"])122deff = lambda w: float(1 + (w.std() / w.mean()) ** 2)123D = {"seed": SEED, "n_productive": int(len(w_prod)), "n_zero": int(len(w_zero)), "n_strata_productive": len(alloc), "floor": FLOOR, "oversample_conf_0.45_0.7": 2.0,124     "population_productive": Ntot, "population_zero": len(zero), "deff_productive": deff(w_prod), "deff_zero": deff(w_zero),125     "half_width_95_precision_p0.5": round(1.96 * math.sqrt(0.25 / len(w_prod) * deff(w_prod)), 4),126     "half_width_95_precision_p0.8": round(1.96 * math.sqrt(0.16 / len(w_prod) * deff(w_prod)), 4),127     "half_width_95_zero_fn_rate": {str(p0): round(1.96 * math.sqrt(p0 * (1 - p0) / len(w_zero) * deff(w_zero)), 4) for p0 in (0.02, 0.05, 0.10)},128     "mde_gate_recall_note": "false-negative rate among zero-mention sections estimated with inverse-probability weights; MDE (80% power, two-sided 5%) for a difference between gate-pass and gate-reject strata of size n1, n2: 2.8*sqrt(p(1-p)*(1/n1+1/n2)*deff)",129     "zero_frame_gate_pass": int(sum(1 for x in sel if x[1] == "zero" and x[14])), "zero_frame_gate_reject": int(sum(1 for x in sel if x[1] == "zero" and not x[14])),130     "split": {"dev": sum(1 for v in split.values() if v == "dev"), "holdout": sum(1 for v in split.values() if v == "holdout")},131     "core_subsample": {"productive": sum(1 for x in sel if x[1] == "productive" and x[0] in core), "zero": sum(1 for x in sel if x[1] == "zero" and x[0] in core),132                        "note": "nested first-pass subsample drawn proportionally to the stratum population within the frame (weights = N_h / n_h,core); the full frame keeps the floor-of-10 design"},133     "rules_column_available": rules.exists(),134     "files_on_node": [str(FR / "frame_dev.jsonl"), str(FR / "frame_holdout.jsonl"), str(FR / "model_output_SEALED.jsonl"), str(FR / "weights.csv")]}135p0 = 0.05136for n1, n2 in ((D["zero_frame_gate_pass"], D["zero_frame_gate_reject"]),):137    D["mde_gate_recall_diff_at_p0.05"] = round(2.8 * math.sqrt(p0 * (1 - p0) * (1 / max(1, n1) + 1 / max(1, n2)) * deff(w_zero)), 4)138r.write_json("design.json", D)139r.dump("frame_composition.csv", con, "SELECT frame, stratum, count(*) n, round(sum(weight)) weighted_n FROM pick GROUP BY 1,2 ORDER BY 1,2")140with open(FR / "README.md", "w") as fh:141    fh.write(f"""# Annotation pilot frames (FIRC-0001, N11) — seed {SEED}, drawn {dt.date.today()}142143frame_dev.jsonl / frame_holdout.jsonl: one JSON object per section with the full section text and `focus_text_seen_by_model`144(exactly what the extractor sent, llm.py:125-126). Annotators label the TEXT; they never see model_output_SEALED.jsonl145until the labels are frozen. weights.csv carries the inverse-probability weight and the split. design.json documents n,146strata, weights, CI half-widths and the MDE. Protocol: docs/annotation_protocol.md in the paper repository.147""")148r.finish()149
Files · 0d1f76183dfb024d676f555cf5ad099a684a4656
Menu

Theme