OCR via Mistral's Batch API - Convert Images into Text at Scale

OCRBatch

Optical Character Recognition (OCR) allows you to retrieve text data from documents. With Mistral OCR, you can do this extremely fast and effectively, extracting text from hundreds and thousands of images (or PDFs).

In this simple cookbook, we will extract text from a set of images using two methods:

  • Without Batch Inference: Looping through the dataset, extracting text from each image, and saving the result.
  • With Batch Inference: Leveraging Batch Inference to extract text with a 50% cost reduction.
    • Inline Batching: Start batching quickly creating a list of payloads and a job directly. Simple and efficient for up to 10k requests.
    • File Batching: Batch by uploading a file of all your payloads, perfect for large workfloads up to 1M requests.

Setup#

First, let's install mistralai, our Python SDK, and datasets, a third party library we will use to download some samples.

!pip install mistralai datasets

We can now set up our client. You can create an API key on AI Studio.

from mistralai.client import Mistral
import getpass
import os

if not os.environ.get("MISTRAL_API_KEY"):
    os.environ["MISTRAL_API_KEY"] = getpass.getpass("Mistral API key: ")

client = Mistral(os.environ["MISTRAL_API_KEY"], timeout_ms=10000)

ocr_model = "mistral-ocr-latest" # For this demo we will use our latest OCR model

Without Batch Inference#

As an example, let's use Mistral OCR to extract text from multiple document images.

We will use a dataset containing image paths, and we will send the images as PDF urls.

For more information on our OCR Processor, please visit our OCR Documentation.

Dataset#

For this demo, we will use a simple dataset containing numerous documents and scans in image format. Specifically, we will use the llamaindex/ParseBench dataset, loaded via the datasets library.

We will download only 100 samples for this demonstration.

from datasets import load_dataset

n_samples = 100
dataset = load_dataset("llamaindex/ParseBench", split="text_content", streaming=True)
subset = list(dataset.take(n_samples))

With our subset of 100 samples ready, we can loop through each image to extract the text.

We will save the results in a new dataset and export it as a JSONL file.

from tqdm import tqdm
import json

ocr_dataset = []
for sample in tqdm(subset):
    image_path = sample['pdf']  # 'image' contains the image path

    # URL of the PDF document
    image_url = f"https://huggingface.co/datasets/llamaindex/ParseBench/resolve/main/{image_path}"

    # Process the image using Mistral OCR
    response = client.ocr.process(
        model=ocr_model,
        document={
            "type": "document_url",
            "document_url": image_url,
        }
    )

    # Store the image data and OCR content in the new dataset
    ocr_dataset.append({
        'image': image_url,
        'ocr_content': "\n\n".join([page.markdown for page in response.pages])
    })

# Saving as a JSON
with open('no_batch_ocr.json', 'w') as f:
    json.dump(ocr_dataset, f, indent=4)

Perfect, we have extracted all text from the 100 samples. However, this process can be made more cost-efficient using Batch Inference.

With Batch Inference#

Batching allows you to run asynchronous inference on large inputs in parallel, reducing compute costs while running large workloads at a 50% discount. We possess two batching approaches:

Inline Batching

For batches of fewer than 10,000 requests, we support inline batching. Instead of creating and uploading a .jsonl file with all the request data (standard approach of File Batching), you can include the request body directly in the job creation request. This is convenient for smaller-scale or less bulk-intensive tasks.

File Batching

The standard batching approach relies on batch files containing all the requests to be processed. We support up to 1 million requests in a single batch, enabling efficient handling of large volumes of requests at a reduced cost. This is ideal for tasks with high throughput requirements but low latency sensitivity or priority.

Inline Batching#

To use Batch Inference with Inline Batching, we need to create a list of batch data with each item corresponding to a request payload. Let's create a function called create_batch_data to handle this task by generating a list in the proper format.

def create_batch_data(image_urls):
    data = []
    for index, url in enumerate(image_urls):
        entry = {
            "custom_id": str(index),
            "body": {
                "document": {
                    "type": "document_url",
                    "document_url": url
                },
                "include_image_base64": True
            }
        }
        data.append(entry)
    return data

The next step involves saving the URL of each image that will be used.

image_urls = []
for sample in tqdm(subset):
    image_path = sample['pdf']  # 'image' contains the image path
    image_url = f"https://huggingface.co/datasets/llamaindex/ParseBench/resolve/main/{image_path}" # URL of the PDF document
    image_urls.append(image_url)

We can now create our batch data.

inline_batch_data = create_batch_data(image_urls)

The data is ready. To initiate the workload, we need to create a job.

inline_created_job = client.batch.jobs.create(
    requests=inline_batch_data,
    model=ocr_model,
    endpoint="/v1/ocr",
    metadata={"job_type": "cookbook"}
)

Our batch job is ready and running!

We can retrieve information such as the status, successful requests or errors:

retrieved_job = client.batch.jobs.get(job_id=inline_created_job.id)
print(f"Status: {retrieved_job.status}")
print(f"Total requests: {retrieved_job.total_requests}")
print(f"Failed requests: {retrieved_job.failed_requests}")
print(f"Successful requests: {retrieved_job.succeeded_requests}")
print(
    f"Percent done: {round((retrieved_job.succeeded_requests + retrieved_job.failed_requests) / retrieved_job.total_requests, 4) * 100}%"
)

Let's automate this feedback loop and download the results once they are ready!

import time
from IPython.display import clear_output

while retrieved_job.status in ["QUEUED", "RUNNING"]:
    retrieved_job = client.batch.jobs.get(job_id=inline_created_job.id)

    clear_output(wait=True)  # Clear the previous output ( User Friendly )
    print(f"Status: {retrieved_job.status}")
    print(f"Total requests: {retrieved_job.total_requests}")
    print(f"Failed requests: {retrieved_job.failed_requests}")
    print(f"Successful requests: {retrieved_job.succeeded_requests}")
    print(
        f"Percent done: {round((retrieved_job.succeeded_requests + retrieved_job.failed_requests) / retrieved_job.total_requests, 4) * 100}%"
    )
    time.sleep(2)

client.files.download(file_id=retrieved_job.output_file) # Download once no longer Queued or Running

File Batching#

To use Batch Inference with File Batching, we need to create a JSONL file containing all the image data and request information for our batch. Let's create a function called create_batch_file to handle this task by generating a file in the proper format.

def create_batch_file(image_urls, output_file):
    data = []
    for index, url in enumerate(image_urls):
        entry = {
            "custom_id": str(index),
            "body": {
                "document": {
                    "type": "document_url",
                    "document_url": url
                },
                "include_image_base64": True
            }
        }
        data.append(entry)
    with open(output_file, 'w') as file:
        file.write('\n'.join([json.dumps(entry) for entry in data]))

The next step involves encoding the data of each image into base64 and saving the URL of each image that will be used.

image_urls = []
for sample in tqdm(subset):
    image_path = sample['pdf']  # 'image' contains the image path
    image_url = f"https://huggingface.co/datasets/llamaindex/ParseBench/resolve/main/{image_path}" # URL of the PDF document
    image_urls.append(image_url)

We can now create our batch file.

batch_file = "batch_file.jsonl"
create_batch_file(image_urls, batch_file)

With everything ready, we can upload it to the API.

batch_data = client.files.upload(
    file={
        "file_name": batch_file,
        "content": open(batch_file, "rb")},
    purpose = "batch"
)

The file is uploaded, but the batch inference has not started yet. To initiate it, we need to create a job.

file_created_job = client.batch.jobs.create(
    input_files=[batch_data.id],
    model=ocr_model,
    endpoint="/v1/ocr",
    metadata={"job_type": "cookbook"}
)

Our batch job is ready and running!

We can retrieve information such as the status, successful requests or errors:

retrieved_job = client.batch.jobs.get(job_id=file_created_job.id)
print(f"Status: {retrieved_job.status}")
print(f"Total requests: {retrieved_job.total_requests}")
print(f"Failed requests: {retrieved_job.failed_requests}")
print(f"Successful requests: {retrieved_job.succeeded_requests}")
print(
    f"Percent done: {round((retrieved_job.succeeded_requests + retrieved_job.failed_requests) / retrieved_job.total_requests, 4) * 100}%"
)

Let's automate this feedback loop and download the results once they are ready!

import time
from IPython.display import clear_output

while retrieved_job.status in ["QUEUED", "RUNNING"]:
    retrieved_job = client.batch.jobs.get(job_id=file_created_job.id)

    clear_output(wait=True)  # Clear the previous output ( User Friendly )
    print(f"Status: {retrieved_job.status}")
    print(f"Total requests: {retrieved_job.total_requests}")
    print(f"Failed requests: {retrieved_job.failed_requests}")
    print(f"Successful requests: {retrieved_job.succeeded_requests}")
    print(
        f"Percent done: {round((retrieved_job.succeeded_requests + retrieved_job.failed_requests) / retrieved_job.total_requests, 4) * 100}%"
    )
    time.sleep(2)

client.files.download(file_id=retrieved_job.output_file) # Download once no longer Queued or Running

Done! With this method, you can now perform OCR tasks in bulk in a very cost-effective way.

If you are interested in data extraction and structured outputs for Document Parsing, we recommend taking a look at our Annotations feature.