spb/firc0001_uqo
Public
1"""N8 - rule-based baseline (snapshot rules.py) vs the LLM output on a stratified seeded sample of ledger2sections, plus rules-on-full vs rules-on-focused text (recall-loss bound, THEO-09). Read-only; CPU only.3"""4from __future__ import annotations5import json, sys, time6from pathlib import Path7import multiprocessing as mp8sys.path.insert(0, str(Path(__file__).resolve().parent))9from common import *10from spid.extract.rules import RuleExtractor11from spid.extract.llm import _focus1213N_PROD = 30000; N_ZERO = 3000014RX = RuleExtractor()1516def work(item):17 sid, name, text = item18 full = RX.extract(text=text, accession_number="a", cik="c", section_id=sid, section_name=name)19 foc = RX.extract(text=_focus(text), accession_number="a", cik="c", section_id=sid, section_name=name)20 return (sid, len(full), sorted({m.project_type for m in full}), len(foc), sorted({m.project_type for m in foc}), len(text),21 [(m.project_type, m.confidence, m.amount_usd, m.description[:200]) for m in full][:20])2223def main():24 r = Run("rules")25 con = connect(SPID_DB)26 tmin = con.execute("select min(extracted_at) from extracted_sections").fetchone()[0]27 con.execute(f"""CREATE TEMP TABLE led AS SELECT e.section_id, e.n_mentions, CASE WHEN e.extracted_at = TIMESTAMP '{tmin}' THEN 'backfilled_june' ELSE 'october' END run,28 f.form_type, m.section_name, m.word_count FROM extracted_sections e JOIN src_filings f ON f.accession_number = e.accession_number29 JOIN (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)""")30 N_prod_pop = con.execute("select count(*) from led where n_mentions>0").fetchone()[0]31 N_zero_pop = con.execute("select count(*) from led where n_mentions=0").fetchone()[0]32 # DuckDB applies USING SAMPLE before WHERE: the stratum filter must sit in a subquery33 con.execute(f"CREATE TEMP TABLE samp AS SELECT *, 'productive' stratum, {N_prod_pop}/{N_PROD}.0 w FROM (SELECT * FROM led WHERE n_mentions>0) USING SAMPLE reservoir({N_PROD} ROWS) REPEATABLE ({SEED})")34 con.execute(f"INSERT INTO samp SELECT *, 'zero', {N_zero_pop}/{N_ZERO}.0 FROM (SELECT * FROM led WHERE n_mentions=0) USING SAMPLE reservoir({N_ZERO} ROWS) REPEATABLE ({SEED})")35 assert con.execute("select count(*) filter (where n_mentions>0), count(*) filter (where n_mentions=0) from samp").fetchone() == (N_PROD, N_ZERO)36 n_s = con.execute("select count(*) from samp").fetchone()[0]37 print(f"sample {n_s} sections (pop productive {N_prod_pop}, zero {N_zero_pop})", flush=True)38 cur = con.execute("""SELECT s.section_id, s.section_name, t.text FROM samp s JOIN (SELECT section_id, text FROM src_sections UNION ALL SELECT section_id, text FROM spid_sections) t USING (section_id)""")39 out = []; t0 = time.time(); done = 040 ctx = mp.get_context("fork")41 with ctx.Pool(4) as pool:42 while True:43 rows = cur.fetchmany(2000)44 if not rows: break45 out.extend(pool.imap_unordered(work, rows, chunksize=25))46 done += len(rows)47 print(f" {done}/{n_s} sections, {time.time()-t0:.0f}s", flush=True)48 import pyarrow as pa, pyarrow.parquet as pq49 tbl = pa.table({"section_id": [o[0] for o in out], "rules_n": [o[1] for o in out], "rules_types": [o[2] for o in out], "rules_focus_n": [o[3] for o in out],50 "rules_focus_types": [o[4] for o in out], "nchar": [o[5] for o in out], "rules_mentions": [json.dumps(o[6]) for o in out]})51 pq.write_table(tbl, r.out / "rules_mentions.parquet") # stays on the node52 con.register("rt", tbl)53 con.execute("""CREATE TEMP TABLE lt AS SELECT section_id, list_distinct(list(project_type)) llm_types FROM project_mentions GROUP BY 1""")54 con.execute("""CREATE TEMP TABLE J AS SELECT s.*, r.rules_n, r.rules_types, r.rules_focus_n, r.rules_focus_types, r.nchar, coalesce(l.llm_types, []) llm_types,55 r.rules_n > 0 rules_pos, s.n_mentions > 0 llm_pos, len(list_intersect(r.rules_types, coalesce(l.llm_types, []))) > 0 type_overlap56 FROM samp s JOIN rt r USING (section_id) LEFT JOIN lt l USING (section_id)""")57 agg = """count(*) n_sampled, round(sum(w)) n_weighted,58 count(*) FILTER (WHERE llm_pos) llm_pos, count(*) FILTER (WHERE rules_pos) rules_pos,59 round(100.0*count(*) FILTER (WHERE llm_pos AND rules_pos)/nullif(count(*) FILTER (WHERE llm_pos),0),2) p_rules_given_llm,60 round(100.0*sum(w) FILTER (WHERE llm_pos AND rules_pos)/nullif(sum(w) FILTER (WHERE rules_pos),0),2) p_llm_given_rules_weighted,61 round(100.0*sum(w) FILTER (WHERE rules_pos)/sum(w),2) rules_pos_pct_weighted, round(100.0*sum(w) FILTER (WHERE llm_pos)/sum(w),2) llm_pos_pct_weighted,62 round(100.0*count(*) FILTER (WHERE llm_pos AND rules_pos AND type_overlap)/nullif(count(*) FILTER (WHERE llm_pos AND rules_pos),0),2) type_overlap_pct_codetected,63 round(100.0*sum(w) FILTER (WHERE llm_pos AND NOT rules_pos)/sum(w),2) llm_only_pct_w, round(100.0*sum(w) FILTER (WHERE rules_pos AND NOT llm_pos)/sum(w),2) rules_only_pct_w,64 round(100.0*sum(w) FILTER (WHERE llm_pos = rules_pos)/sum(w),2) agreement_pct_w,65 round(avg(rules_n) FILTER (WHERE llm_pos),3) rules_mentions_per_llm_pos_section, round(avg(rules_n) FILTER (WHERE NOT llm_pos),3) rules_mentions_per_llm_zero_section"""66 r.dump("agreement_by_form_item.csv", con, f"SELECT form_type, section_name item, {agg} FROM J GROUP BY 1,2 ORDER BY 1,2")67 r.dump("agreement_by_form.csv", con, f"SELECT form_type, {agg} FROM J GROUP BY 1 ORDER BY 1")68 r.dump("agreement_by_run.csv", con, f"SELECT run, {agg} FROM J GROUP BY 1 ORDER BY 1")69 r.dump("agreement_overall.csv", con, f"SELECT 'all' scope, {agg} FROM J")70 r.dump("type_agreement.csv", con, """71 WITH x AS (SELECT unnest(llm_types) t, rules_types FROM J WHERE llm_pos AND rules_pos)72 SELECT t llm_type, count(*) n, count(*) FILTER (WHERE list_contains(rules_types, t)) same_type_in_rules, round(100.0*count(*) FILTER (WHERE list_contains(rules_types, t))/count(*),2) pct FROM x GROUP BY 1 ORDER BY 2 DESC""")73 r.dump("focus_recall_loss.csv", con, """74 SELECT form_type, section_name item, count(*) FILTER (WHERE nchar > 9000) n_over_budget,75 sum(rules_n) FILTER (WHERE nchar > 9000) rules_full, sum(rules_focus_n) FILTER (WHERE nchar > 9000) rules_focused,76 round(100.0*(1 - sum(rules_focus_n) FILTER (WHERE nchar > 9000)*1.0/nullif(sum(rules_n) FILTER (WHERE nchar > 9000),0)),2) pct_rules_mentions_lost_to_focus,77 round(100.0*count(*) FILTER (WHERE nchar > 9000 AND rules_n > 0 AND rules_focus_n = 0)/nullif(count(*) FILTER (WHERE nchar > 9000 AND rules_n > 0),0),2) pct_sections_lost_entirely,78 round(100.0*sum(w) FILTER (WHERE nchar > 9000 AND rules_n > 0 AND rules_focus_n = 0)/nullif(sum(w) FILTER (WHERE nchar > 9000 AND rules_n > 0),0),2) pct_sections_lost_entirely_w79 FROM J GROUP BY 1,2 ORDER BY 1,2""")80 r.dump("focus_recall_loss_overall.csv", con, """81 SELECT 'over_9000_chars' scope, count(*) n, sum(rules_n) rules_full, sum(rules_focus_n) rules_focused, round(100.0*(1 - sum(rules_focus_n)*1.0/nullif(sum(rules_n),0)),2) pct_lost,82 round(100.0*sum(w*rules_n)/nullif(sum(w*rules_n) + 0,0),2) dummy FROM J WHERE nchar > 9000""")83 r.dump("disagreement_strata.csv", con, """SELECT form_type, section_name item, run, llm_pos, rules_pos, count(*) n_sampled, round(sum(w)) n_weighted FROM J GROUP BY 1,2,3,4,5 ORDER BY 1,2,3,4,5""")84 r.notes["sample"] = {"productive": N_PROD, "zero": N_ZERO, "pop_productive": N_prod_pop, "pop_zero": N_zero_pop, "seed": SEED}85 r.finish()8687if __name__ == "__main__":88 main()89