Prompt Optimizer for Structured Document Extraction

Automatically improves the extraction prompt for any PDF→JSON schema using greedy beam-search. Point it at a new schema, run it, get a better prompt — no code changes.

Built on ExtractBench (MIT).

Demo video: (https://www.youtube.com/watch?v=0uPxbf7_0Gk)

---

How it works

flowchart TD
    A[PDFs + Gold JSON] --> B[Extractor<br/>Gemini PDF → JSON]
    B --> C[Scorer<br/>F1 + Semantic Metrics]
    C --> D[Optimizer<br/>Greedy Beam Search]
    D --> E[Proposer<br/>Groq Llama-3]
    E --> F[Improved Prompt]
    F --> G[REPORT.md + Logs]

Loading

---

Quickstart

# 1. Clone and install
git clone <this-repo> && cd pibit
pip install -r requirements.txt

# 2. Add API keys (.env or environment)
echo "GEMINI_API_KEY=your_key" >> .env
echo "GROQ_API_KEY=your_key"   >> .env

# 3. Get the dataset
git clone https://github.com/ContextualAI/extract-bench data/extract-bench

# 4. Run
python main.py run --config config/hiring_resume.yaml

---

A run in action

---

Results snapshot

---

Available configs

Config file	Schema	Notes
config/hiring_resume.yaml	hiring/resume	fast, good for testing
config/academic_research.yaml	academic/research
config/finance_10kq.yaml	finance/10kq	chunked extraction (large output)
config/finance_credit_agreement.yaml	finance/credit_agreement
config/sport_sport.yaml	sport/swimming

To retarget to any schema: only edit the config file. No code changes.

---

Inspecting a run

# Score trajectory across iterations
python inspect_run.py trajectory runs/<run_id>

# Diff prompts between two iterations
python inspect_run.py diff runs/<run_id> 0 5

# Per-field F1 breakdown  — find where the model is losing points
python inspect_run.py fields runs/<run_id>

# Cost and latency per LLM call
python inspect_run.py costs runs/<run_id>

# List all runs
python inspect_run.py list

---

Resuming an interrupted run

Every iteration is written to disk before the next one starts. If the process dies mid-run:

python main.py resume --run-dir runs/<run_id>

It picks up from the last completed iteration — no re-work, no re-spending.

---

Run directory layout

runs/<run_id>/
  config.json          snapshot of config at run time
  iteration_000.json   seed evaluation (prompt, val scores, per-field breakdown)
  iteration_001.json   …and so on
  best_prompt.txt      current best prompt, updated each iteration
  llm_calls.jsonl      every LLM call: input · output · cost · latency
  semantic_cache/      cached stochastic scorer results (keyed by content hash)
  final.json           seed + final test scores
  REPORT.md            generated report

---

Optimization algorithm

Greedy beam search — close to Karpathy's autoresearch loop, extended to support a beam of width > 1 and three mutation strategies:

1. Evaluate seed prompt on validation split → initial beam.
2. Each iteration:
   • Proposer generates K mutations per beam member.
   • Quick-score all mutations on a mini-batch.
   • Full-score top survivors on the whole validation split.
   • Keep the best beam_size prompts by F1; accept if any beats the current best.
3. On stall (N iterations with no improvement), escalate: targeted → few_shot → exploratory.
4. On budget exhaustion, evaluate the current best on the held-out test split.

Pathological cases:

Situation	What happens
Mutation is worse	Rejected; beam unchanged
Duplicate proposal	Detected by prompt hash; silently skipped
N iterations with no gain	Strategy escalation (see above)
Dollar / iteration cap reached	Loop exits; best prompt is tested

---

Scoring

Each field in the schema carries an evaluation_config. The scorer honours it exactly:

Config type	Method
string_exact	Case-insensitive, whitespace-normalised exact match
string_semantic	LLM judge — result cached per SHA-256(predicted, gold)
integer_exact	Exact integer comparison with type coercion
number_tolerance	Within ±tolerance% of gold value (default 5%)
array_llm	Per-element LLM judge for precision and recall (cached)

Array alignment policy: greedy best-match on a pairwise similarity matrix. Pairs with similarity ≥ 0.5 count as true positives; unmatched predictions → false positives; unmatched gold → false negatives. Object similarity = mean sub-field F1; scalar similarity = scalar metric score.

Stochastic metrics (string_semantic, array_llm) are cached per content hash in runs/<run_id>/semantic_cache/ — deterministic across runs for the same (predicted, gold) pair.

The scorer is a standalone module with no dependency on the optimizer. It can be used independently.

---

Data split policy

Splits are deterministic and seeded. The seed and ratios are set per config:

dataset:
  split_seed:  42
  train_ratio: 0.6
  val_ratio:   0.2
  # remainder is test

The same seed always produces the same split. No documents are shared across splits. No external data is added.

---

Project structure

main.py              CLI — run / resume commands
inspect_run.py       Observability tooling
src/
  dataset.py         ExtractBench loader + deterministic split
  extractor.py       PDF → JSON via Gemini (inline PDF + chunked mode)
  scorer.py          Per-field P/R/F1 following evaluation_config
  proposer.py        Prompt mutation proposals via LLM meta-prompting
  optimizer.py       Beam-search loop
  llm_client.py      Gemini API wrapper — cost tracking, caching, retry
  groq_client.py     Groq wrapper for fast proposer calls
  persistence.py     State save/load + run resumption
  report.py          REPORT.md generation
config/              Per-schema YAML configs
tests/               Scorer unit tests (39 tests)
docs/                Screenshots used in this README

---

Running tests

pytest tests/ -v
# 39 passed

---

Configuration reference

Key parameters — full list is in any config file:

llm:
  extractor_model:      gemini-2.5-flash       # must support PDF input
  extractor_max_tokens: 8192
  proposer_provider:    groq                   # or gemini
  proposer_model:       llama-3.1-8b-instant
  scorer_model:         gemini-2.5-flash-lite

extraction:
  extraction_cache_enabled: true
  chunk_by_section: false   # set true for large-output schemas (10kq)

optimization:
  beam_size:               1
  mutations_per_iteration: 3
  max_iterations:          10
  max_cost_usd:            5.0
  stall_threshold:         3
  initial_strategy:        targeted   # targeted | few_shot | exploratory