How to improve ai_extract accuracy, version prompts and schemas, evaluate results, and how Genie Code can automate them for you
Document extraction often starts with a simple loop: spot an incorrect field, adjust the prompt or schema, rerun the document, and check the result.
That approach becomes unreliable as fields and document formats multiply. A change may fix one field while introducing regressions elsewhere, work for one layout but fail on another, or appear better simply due to normal variation across LLM runs.
We saw this on a recent project that extracted roughly 40 structured fields from partner-submitted order documents across dozens of PDF layouts. Reviewers corrected and approved extracted values before they moved downstream, giving us human-reviewed ground truth for evaluation.
The challenge was not getting ai_extract
to return the fields. It was improving extraction without degrading the system overall.
For example, the revision
field sometimes returned status terms such as ORIGINAL
instead of a revision
number or blank. A small schema change cut errors for that field roughly in half. But across every field, total extraction errors increased by 9%.
The lesson was clear: improving one field does not necessarily improve the extraction system.
The pipeline already had most of the pieces needed to evaluate changes systematically. Documents landed in object storage and flowed through a Lakeflow SDP pipeline using Auto Loader. Bronze parsed documents with ai_parse_document, Silver extracted structured fields with ai_extract, and Gold published results as a materialized view. Predictions synced to Lakebase, where a Databricks App let reviewers correct and approve fields, and Lakebase CDC returned those approvals to the lakehouse.
Two pieces were missing from the feedback loop. Reviewer approvals were available, but they were not assembled into a reusable evaluation dataset. At the same time, prompts and schemas were embedded directly in pipeline code, so even a small change to extraction required a code deployment.
We addressed both problems by building a workflow around three questions: What should we change? How should we version it? And what evidence is enough to promote it?
This post covers three parts of that workflow:
Together, these steps create a repeatable path from a proposed extraction change to production, with evidence that the change improves the system as a whole.
The rest of this post walks through the practices we recommend and why each one matters. You don't have to implement all of them on your own. Genie Code for document processing builds the Lakeflow job (ai_parse_document, ai_classify, and ai_extract with citations and confidence scores) and automates the core of this workflow with an evaluation and hill-climbing loop:
For example, you can ask Genie Code: "Evaluate the extraction against my labeled table and tell me which fields are weakest," followed by "Accuracy on the revision field could be higher. Iteratively improve the schema and instructions, re-run the evaluation, and show me the before and after."
Understanding the fundamentals still matters. It helps you judge Genie Code's proposed changes, decide when a result is good enough to promote, and extend the workflow for production, for example by versioning prompts and schemas in MLflow Prompt Registry. Each section below explains a recommended practice and notes how Genie Code can handle it for you.
ai_extract
leverai_extract
takes three inputs: the parsed document, a JSON schema, and optional global instructions.
The schema and instructions serve different purposes:
The result is a structured value with one entry for each schema field. You can access individual fields directly, for example: extracted:response.revision.value
.
Start with the schema structure. When a field has a closed set of valid values, use an enum constraint rather than listing the allowed values in the description. Use numeric types for numbers and integers, and use nested arrays or objects for repeated structures.
Then use schema descriptions for rules that types alone cannot express, such as:
004
ORIGINAL
, NEW
, or REVISED
null
when a value isn't presentKEY: Keeping field rules in the schema makes the extraction contract easier to inspect and prevents the global prompt from becoming a list of unrelated exceptions.
Reserve the global prompt for rules that apply across the document, including:
For example, if entity names are mistakenly extracted from boilerplate footers or third-party contact blocks, a FORBIDDEN-SOURCES
instruction can tell the model to ignore those regions.
A useful rule is: If the rule defines one field, put it in the schema. If it defines how to interpret the document, put it in the prompt.
Automate it with Genie Code: The hill-climbing loop follows the same principle. Its proposed fixes most often sharpen a field description in the schema, and are shown as a diff for you to approve before they apply.
Not every transformation needs an LLM. Handle casing, separator removal, zero-padding, and mappings such as TELEVISION to TV with SQL or another deterministic transformation when possible.
KEY: Use the model to interpret the document; use code to normalize the output.
Sometimes the instructions are clear, but the document is difficult to follow. Large schemas, repeated line items, complex tables, and multi-page layouts can make it hard to keep values aligned with the right rows.
For these cases, ai_extract
version 2.1 offers precision mode:
Precision Mode brings frontier extraction capabilities for more demanding use cases. With the prompt and schema held constant, it tends to help most on repeated structures and reduces over-extraction (the model returning a value where the correct answer is blank).
Diagnose the failure before changing the configuration:
The approaches are complementary. In our testing, the schema rule corrected a field the model misunderstood, while precision mode improved extraction of fields from complex tables.
When an extracted value is wrong, citations help show where the model found its answer:
Citations provide field-level bounding boxes that can be mapped back to element spans in the source document. This helps distinguish between two different types of errors:
These require different fixes. Extraction errors may call for a change to the schema or prompt. Labeling disagreements should be resolved by defining the expected convention in the reviewer guidelines.
Confidence scores can also help prioritize human review:
Use them as a ranking signal, such as reviewing low-confidence fields first, rather than as an automatic quality gate. A high-confidence result can still be wrong.
In Genie Code: The extraction tasks Genie Code builds include citations and confidence scores. You can ask it to "Set confidence thresholds per field from the evaluation results and route low-confidence rows to review," and use the Document Parsing UI to visually debug parsing and extraction.
Reduce the friction of hardcoding prompts and schemas in notebooks by managing them as configurable, reusable definitions outside the pipeline code. A small wording change requires a code deployment, previous versions are difficult to reconstruct, and there is no evidence that the new version is better than the old one.
MLflow Prompt Registry lets you version those inputs separately from pipeline code. Version both the prompt and the schema: field descriptions influence extraction just as global instructions do. Version only the prompt, and half of your extraction logic is still traveling as an untracked code change.
Beyond the Genie Code loop: Genie Code helps you find the right prompt and schema. When you take that configuration to production, versioning it lets you promote and roll back changes without a code deployment.
A simple alias strategy might use:
@prod
for the currently approved version@candidate
for the version under evaluation@scr_*
for experimental versions being compared in parallelThe pipeline then loads the active versions by alias:
The resolved prompt and schema are then passed to ai_extract
.
When a candidate is approved, mlflow.genai.set_prompt_alias
can move @prod
to the new version. The pipeline resolves it on the next run without changing the pipeline code. To roll back, restore both aliases to the previous recorded versions. The underlying versions remain in the registry.
This separates extraction configuration from the code that runs it. Figure 2 also adds the other missing piece: reusable ground truth.
The ground-truth view pairs approved labels with source documents and predictions, creating the reference dataset for measuring changes.
Version history tells you what changed. It doesn't tell you whether the change was better. Build a fixed cohort from reviewer-approved results; ours was several hundred PDFs covering a wide range of layouts. Score the production prompt and each candidate prompt against the same test set.
For every version, track metrics such as:
Keep both the documents and metric definitions consistent across comparisons. Otherwise, a score change could reflect a different test set or scoring rule rather than a different extraction configuration.
In Genie Code: Point Genie Code at your labeled table, such as the ground-truth view in Figure 2, and ask it to evaluate the extraction. It scores each field and re-runs the same eval set after each accepted change, so before-and-after comparisons use the same documents. The noise-floor, coverage, and verdict checks below are useful guardrails to add when you take a change to production.
A candidate scoring slightly higher once doesn't mean it is better. Run the production configuration at least 3-5 times on the frozen test set and calculate its mean accuracy and standard deviation (σ). That σ is your noise floor. Your policy might require a candidate's improvement to exceed 2σ before treating it as meaningful; anything inside that band is inconclusive rather than a win.
Accuracy can also be misleading if you ignore missing predictions.
Suppose you score only fields where both the model and reviewer supplied a value. A candidate that leaves difficult fields blank can appear more accurate because those misses disappear from the denominator. The reverse can also happen. When we scored the revision
field only on documents where the reviewer supplied a revision
value, our fix appeared to improve accuracy by less than 3 percentage points. But when we counted all errors, including cases where the correct answer was blank, the fix cut revision
errors roughly in half. Most of the improvement came from correctly returning no value where one did not belong.
That is why coverage should be reported alongside accuracy. Define it explicitly, for example, as the share of non-null ground-truth fields for which the model returns a prediction. If accuracy rises but coverage falls materially, investigate before promoting the candidate.
The resulting harness follows a repeatable sequence: freeze the test set, establish the production baseline, register a candidate on a screening alias, run repeated evaluations, and compare accuracy, coverage, and stability.
Once you measure accuracy, coverage, and run-to-run variation, your evaluation harness can reduce those metrics to a small set of outcomes:
These verdicts provide a consistent way to summarize an evaluation before deciding whether the candidate should be promoted to production.
We scored each candidate three times on the same test set, counting errors in every field, relative to the original prompt and schema.
Change | revision errors | Errors across all fields | Verdict |
---|---|---|---|
revision schema fix only | -49% | +9% | TARGET_IMPROVED_WITH_REGRESSION |
Restructured prompt and schema | -57% | -16% | BETTER |
Precision mode only, no prompt or schema change | -10% | -24% | BETTER |
Restructured prompt and schema + precision mode | -64% | -32% | BETTER |
The spot fix did what we wanted for revision
, but the model started inventing values in unrelated line-item fields such as schedule details. Editing one field's description is not a one-field experiment: the model reads the schema as a whole.
So we did not ship it.
Instead, we moved field-specific rules out of the global prompt and into schema descriptions, including the same revision
rule. The prompt ended up at less than half its original length, and errors fell by 16% overall, with estimate
errors down by 30%.
Two fields worsened, and they went on the target list for the next round.
Precision mode told a different story. With no prompt or schema change, it cut errors by 24%, mostly by no longer inventing values in repeated line-item fields. But it barely changed revision
errors because that was a definition problem, and only the schema rule fixed it.
Combined, they stack: precision mode on top of the restructured prompt and schema removed another 19% of the remaining errors. Overall, errors fell by about a third (32%), and revision
errors by nearly two-thirds (64%).
This is the regression Genie Code's loop is designed to surface. Because every accepted change re-runs the full eval set, a spot fix that helps revision but hurts line-item fields appears in the before-and-after results rather than in production.
A verdict is a starting point, not the final decision.
One earlier candidate cut estimate errors by about a third and was only INCONCLUSIVE overall. But it had also started filling in a yes/no flag on nearly every document, while reviewers had marked that field on only a handful of documents.
A severe regression in one field can preclude promotion even when aggregate metrics appear acceptable. Promotion should therefore remain human-controlled.
If a candidate is classified as BETTER
, an owner can update the @prod
alias and record the candidate version, evaluation metrics, author, and timestamp.
There may also be valid reasons to promote an INCONCLUSIVE
change, such as implementing an important business convention not fully captured by aggregate accuracy. In those cases, recording an override reason preserves the rationale behind the decision.
With that review and alias update, the feedback loop closes:
Reliable document extraction is not about finding one perfect prompt. It is about choosing the appropriate change, tracking what changed, and testing whether the result improved.
With ai_extract
, that means using schema descriptions for field-level behavior, global instructions for document-level reasoning, precision mode when structure is the bottleneck, and citations when you need to understand why an extraction failed.
MLflow Prompt Registry then gives you a versioned path for managing prompts and schemas, while a fixed evaluation test set gives you evidence for deciding what should reach production.
Together, those pieces turn prompt iteration from a series of one-off edits into a repeatable engineering workflow.
Genie Code for document processing automates the core of this loop (measure, improve, re-evaluate), so most teams can start there, then add Prompt Registry versioning, noise-floor checks, and promotion gates as the pipeline moves to production.
Explore the Databricks AI Functions documentation to learn more about ai_extract
, and see the MLflow Prompt Registry documentation for prompt versioning and alias management. Try Genie Code for document processing. Open Genie Code from Jobs and Pipelines, point it at your documents and a labeled table, and ask it to evaluate and improve your extraction.
Solutions, Customers, Best Practices, Media & Entertainment and Data Science and ML
Subscribe to our blog and get the latest posts delivered to your inbox.