DEV Community

Cover image for Create the fastest and precise invoice data extractor for structural output using AI
Wilbert Misingo
Wilbert Misingo

Posted on • Edited on

Create the fastest and precise invoice data extractor for structural output using AI

Using LlamaExtract with Pydantic Models for Shop Receipts Extraction

In this article, we'll explore how to use LlamaExtract incorporated with schemas from Pydantic models inorder to extract structured data from shop receipts. This approach helps in organizing receipt information systematically, making it easier to analyze and manage.

Setup

First, ensure you have the llama-extract client library installed. Use the following command:

pip install llama-extract pydantic
Enter fullscreen mode Exit fullscreen mode

Note: If you see a notice about updating pip, you may update it using the command provided.

First, login and get an api-key for free from Llama Index Cloud

Set up the environment variable for your LlamaExtract API key:

import os

os.environ["LLAMA_CLOUD_API_KEY"] = "YOUR LLAMA INDEX CLOUD API HERE"
Enter fullscreen mode Exit fullscreen mode

Load Data

For this example, let's assume we have a dataset of shop receipts in PDF format. Place these files in a directory named receipts.

DATA_DIR = "data/receipts"
fnames = os.listdir(DATA_DIR)
fnames = [fname for fname in fnames if fname.endswith(".pdf")]
fpaths = [os.path.join(DATA_DIR, fname) for fname in fnames]
fpaths
Enter fullscreen mode Exit fullscreen mode

The output should list the file paths of the receipts:

['data/receipts/receipt.pdf']
Enter fullscreen mode Exit fullscreen mode

Define a Pydantic Model

We'll define our data model using Pydantic, this would tell the API which fields/data we are expecting or want to extract from the PDF. For shop receipts, we might be interested in extracting the store name, date, total amount, and list of items purchased.

from pydantic import BaseModel
from typing import List

class Item(BaseModel):
    name: str
    quantity: int
    price: float

class Receipt(BaseModel):
    store_name: str
    date: str
    total_amount: float
    items: List[Item]
Enter fullscreen mode Exit fullscreen mode

Create Schema

Now, we can use the Pydantic model to define an extraction schema in LlamaExtract.

from llama_extract import LlamaExtract

extractor = LlamaExtract(verbose=True)
schema_response = await extractor.acreate_schema("Receipt Schema", data_schema=Receipt)
schema_response.data_schema
Enter fullscreen mode Exit fullscreen mode

The output schema should resemble the following:

{
    'type': 'object',
    '$defs': {
        'Item': {
            'type': 'object',
            'title': 'Item',
            'required': ['name', 'quantity', 'price'],
            'properties': {
                'name': {'type': 'string', 'title': 'Name'},
                'quantity': {'type': 'integer', 'title': 'Quantity'},
                'price': {'type': 'number', 'title': 'Price'}
            }
        }
    },
    'title': 'Receipt',
    'required': ['store_name', 'date', 'total_amount', 'items'],
    'properties': {
        'store_name': {'type': 'string', 'title': 'Store Name'},
        'date': {'type': 'string', 'title': 'Date'},
        'total_amount': {'type': 'number', 'title': 'Total Amount'},
        'items': {
            'type': 'array',
            'title': 'Items',
            'items': {'$ref': '#/$defs/Item'}
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

Run Extraction

With the schema defined, we can now extract structured data from our receipt files. By specifying Receipt as the response model, we ensure the extracted data is validated and structured.

responses = await extractor.aextract(
    schema_response.id, fpaths, response_model=Receipt
)

Enter fullscreen mode Exit fullscreen mode

You can access the raw JSON output if needed:

data = responses[0].data
print(data)
Enter fullscreen mode Exit fullscreen mode

Example JSON output:

{
    'store_name': 'ABC Electronics',
    'date': '2024-08-05',
    'total_amount': 123.45,
    'items': [
        {'name': 'Laptop', 'quantity': 1, 'price': 999.99},
        {'name': 'Mouse', 'quantity': 1, 'price': 25.00},
        {'name': 'Keyboard', 'quantity': 1, 'price': 50.00}
    ]
}
Enter fullscreen mode Exit fullscreen mode

Conclusion

In this article, we demonstrated how to use LlamaExtract with Pydantic models for defining data schemas and extract structured data from shop receipts. This approach ensures that the extracted information is well-organized and validated, making it easier to handle and analyze.

This can also be use for many cases, invoices, receipts, reports etc.

Happy Coding!!

Do you have a project πŸš€ that you want me to assist you email me🀝😊: wilbertmisingo@gmail.com

Have a question or wanna be the first to know about my posts:-
Follow βœ… me on GitHub
Follow βœ… me on LinkedIn πŸ’Ό
Follow βœ… me on Twitter/X 𝕏

Top comments (0)