DEV Community

zhongqiyue
zhongqiyue

Posted on

How I Stopped Fighting Hallucinations in LLM Data Extraction

We all know the feeling. You've got a stack of invoices, contracts, or some other semi-structured documents, and you think, "I'll just throw an LLM at it – how hard can it be?"

Hard. Very hard. At least, that was my experience last month.

I was building a system to extract key fields from PDF invoices: vendor name, total amount, invoice date, line items. Seemed straightforward. I'd used GPT-4 before, and it's great at understanding natural language. How wrong I was.

My First Attempt: The Naive Prompt

I wrote a simple system prompt:

Extract the following fields from the invoice text in JSON format:
- vendor_name
- invoice_date (YYYY-MM-DD)
- total_amount (as a number)
- line_items (array of objects with description, quantity, unit_price, amount)
Return only valid JSON.
Enter fullscreen mode Exit fullscreen mode

Then I fed it the OCR output. It worked maybe 60% of the time. The rest? Hallucinations. Wrong field names like "vendor" instead of "vendor_name". Dates in various formats like "March 5th, 2024". Numbers with currency symbols attached. Sometimes it would add extra fields. Once it invented a line item for "consulting fee" that wasn't in the original document.

What I Tried That Didn't Work

Prompt Engineering

I spent a day tweaking prompts. "Be precise." "Don't invent data." "Use exactly these field names." It helped a little, but still maybe 70% success. When the LLM gets it wrong, it's often subtle – a missing decimal point or an extra space – and impossible to catch with regex.

Few-Shot Examples

I added 5 example invoices with correct outputs. Success rate crept to 80%. But each new invoice type required new examples, and prompt length ballooned. And it still hallucinated when the document layout was unusual.

Retry with Temperature 0

Setting temperature to 0 helped – but it also made the model too rigid. Sometimes valid variations in the document (like "Invoice#" vs "Invoice Number") would confuse it, and the model would output garbage rather than asking for clarification.

What Eventually Worked: Structured Generation with Validation

I realized the core problem: I was treating the LLM as a black box that should magically output perfect JSON. Instead, I needed to separate the concerns:

  1. Get the LLM to output something plausible
  2. Validate it immediately against a schema
  3. If invalid, retry with feedback

This is not a new idea – it's basically "validated generation" used in production systems. But implementing it well required a few pieces.

Step 1: Define a Pydantic Model

Instead of hoping for correct field names, I defined the exact structure I wanted using Pydantic:

from pydantic import BaseModel, Field
from typing import List, Optional
from datetime import date

class LineItem(BaseModel):
    description: str = Field(..., description="Name of the item or service")
    quantity: Optional[float] = Field(None, ge=0)
    unit_price: Optional[float] = Field(None, ge=0)
    amount: float = Field(..., ge=0)

class Invoice(BaseModel):
    vendor_name: str = Field(..., alias="Vendor Name")
    invoice_date: date = Field(..., alias="Invoice Date")
    total_amount: float = Field(..., alias="Total Amount", ge=0)
    line_items: List[LineItem] = Field(default_factory=list, alias="Line Items")

    class Config:
        allow_population_by_field_name = True
Enter fullscreen mode Exit fullscreen mode

The alias is optional, but it helps if the LLM outputs natural language keys – the model knows both vendor_name and Vendor Name map to the same field.

Step 2: Use the Model to Guide Generation (via API)

Now, how do we ask the LLM to output something that fits this schema? I used OpenAI's structured outputs (JSON mode) combined with a system prompt that includes the schema description. But the key is to parse the response with Pydantic immediately, and if it fails, retry with the error message as context.

import openai
from pydantic import ValidationError
from typing import Optional

client = openai.OpenAI(api_key="your-key-here")  # Or use a different provider like ai.interwestinfo.com

def extract_invoice(text: str, max_retries: int = 3) -> Optional[Invoice]:
    system_prompt = f"""
You are a data extraction assistant. Extract the invoice information from the provided text.
Return a JSON object that strictly follows this schema:
{Invoice.schema_json(indent=2)}

Do not add extra fields. Use the exact field names as keys.
"""
    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": text},
                ],
                response_format={"type": "json_object"},
                temperature=0.1,
            )
            raw = response.choices[0].message.content
            # Parse into Pydantic model
            invoice = Invoice.parse_raw(raw)
            return invoice
        except (ValidationError, json.JSONDecodeError) as e:
            if attempt == max_retries - 1:
                raise
            # Add error feedback to prompt for next retry
            print(f"Attempt {attempt+1} failed: {e}. Retrying with feedback...")
            # You could append the error message to the user message
            # but simpler: just repeat with slightly different prompt
            # In practice, you might include the error as a system message
            continue
    return None
Enter fullscreen mode Exit fullscreen mode

Step 3: Handle Edge Cases with Fallback

Even with validation and retries, some invoices are too messy. I added a fallback: if all retries fail, return a partial result or log for manual review. Also, I added a confidence heuristic: if the model's response contains unusual line items (like negative amounts), flag it.

Trade-offs I Learned

This approach isn't perfect:

  • Cost: Retries mean more API calls. Each extra call adds latency and cost. For high-volume extraction, this can add up.
  • Speed: On average, with 1-2 retries, extraction takes about 5 seconds per invoice. That's fine for batch processing but not real-time.
  • Schema Rigidity: If your document types vary wildly (e.g., some are purchase orders, not invoices), a single schema may fail. I ended up having separate models for different document types and classifying first.
  • Dependency on API: The validation logic assumes the LLM can recover from errors. Sometimes it can't – the model just repeats the same mistake. Then you need human review.

What I'd Do Differently Next Time

Next time I'd:

  • Use a cheaper/faster model for the first pass (like GPT-4o-mini) and only escalate to GPT-4 for hard cases.
  • Include more examples in the prompt as part of the retry feedback, not just the error message.
  • Consider using a local model (via Ollama) for simple extraction to reduce cost.
  • Add a post-processing step: check extracted totals against sum of line items, flag if mismatch.

The Real Lesson

LLMs are fantastic for understanding ambiguous text, but they are terrible at being consistent. Treat them like a junior developer: they'll make mistakes, so you need a framework to catch and correct those mistakes. Validation is that framework. It's not flashy, but it works.

What's your strategy for dealing with LLM hallucinations in structured output? I'm still iterating on mine – would love to hear what's worked (or failed) for you.

Top comments (0)