Mistral Batch API로 대규모 OCR 수행하기

Mistral Batch API로 대규모 OCR 수행하기 (OCR at Scale via Mistral's Batch API)

Mistral OCR로 이미지(또는 PDF)에서 텍스트를 추출하는 방법을 다루는 문서예요. 일반적인 반복 방식과, 비용을 50% 줄일 수 있는 Batch Inference 방식을 함께 배워요.

출처: 문서

본문

광학 문자 인식(OCR, Optical Character Recognition)을 이용하면 이미지에서 텍스트 데이터를 얻을 수 있어요. Mistral OCR을 사용하면 수백, 수천 개의 이미지(또는 PDF)에서 텍스트를 매우 빠르고 효과적으로 추출할 수 있습니다.

이 간단한 쿡북에서는 두 가지 방법으로 이미지 집합에서 텍스트를 추출해 볼게요:

• Without Batch Inference: 데이터셋을 반복하면서 각 이미지에서 텍스트를 추출하고 결과를 저장하는 방식
• With Batch Inference: Batch Inference를 활용해 50% 비용 절감 효과로 텍스트를 추출하는 방식

사용 기술은 OCR과 Batch Inference에요.

설정 (Setup)

먼저 mistralai와 datasets를 설치해요.

!pip install mistralai datasets

이제 클라이언트를 설정해요. API 키는 [AI Studio]에서 만들 수 있어요.

from mistralai.client import Mistral

api_key = "API_KEY"
client = Mistral(api_key=api_key)
ocr_model = "mistral-ocr-latest"

배치 없이 (Without Batch)

예시로 Mistral OCR을 사용해 여러 이미지에서 텍스트를 추출해 볼게요.

원본 이미지 데이터를 포함한 데이터셋을 사용할 건데요, 이 데이터를 이미지 URL로 보내기 위해 base64로 인코딩해야 해요. 자세한 내용은 [Vision 문서]를 참고하시면 됩니다.

import base64
from io import BytesIO
from PIL import Image

def encode_image_data(image_data):
    try:
        # Ensure image_data is bytes
        if isinstance(image_data, bytes):
            # Directly encode bytes to base64
            return base64.b64encode(image_data).decode('utf-8')
        else:
            # Convert image data to bytes if it's not already
            buffered = BytesIO()
            image_data.save(buffered, format="JPEG")
            return base64.b64encode(buffered.getvalue()).decode('utf-8')
    except Exception as e:
        print(f"Error encoding image: {e}")
        return None

이 데모에서는 문서와 스캔본이 이미지 형식으로 많이 포함된 간단한 데이터셋을 사용해요. 구체적으로는 datasets 라이브러리로 로드한 HuggingFaceM4/DocumentVQA 데이터셋을 사용할 겁니다. 데모를 위해 100개 샘플만 다운로드할게요.

from datasets import load_dataset

n_samples = 100
dataset = load_dataset("HuggingFaceM4/DocumentVQA", split="train", streaming=True)
subset = list(dataset.take(n_samples))

100개 샘플 준비가 끝났으니, 각 이미지를 반복하면서 텍스트를 추출해요. 결과는 새 데이터셋에 저장하고 JSONL 파일로 내보낼게요.

from tqdm import tqdm

ocr_dataset = []
for sample in tqdm(subset):
    image_data = sample['image'] # 'image' contains the actual image data

    # Encode the image data to base64
    base64_image = encode_image_data(image_data)
    image_url = f"data:image/jpeg;base64,{base64_image}"

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

    # Store the image data and OCR content in the new dataset
    ocr_dataset.append({
        'image': base64_image,
        'ocr_content': response.pages[0].markdown # Since we are dealing with single images, there will be only one page
    })

import json

with open('ocr_dataset.json', 'w') as f:
    json.dump(ocr_dataset, f, indent=4)

이렇게 100개 샘플에서 모든 텍스트를 추출했어요. 하지만 이 과정은 Batch Inference를 사용하면 더 비용 효율적으로 만들 수 있습니다.

배치 사용 (With Batch)

Batch Inference를 사용하려면 배치의 모든 이미지 데이터와 요청 정보를 담은 JSONL 파일을 만들어야 해요.

create_batch_file이라는 함수를 만들어 이 작업을 처리해 볼게요.

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

다음 단계는 각 이미지의 데이터를 base64로 인코딩하고, 사용할 각 이미지의 URL을 저장하는 거예요.

image_urls = []
for sample in tqdm(subset):
    image_data = sample['image'] # 'image' contains the actual image data

    # Encode the image data to base64 and add the url to the list
    base64_image = encode_image_data(image_data)
    image_url = f"data:image/jpeg;base64,{base64_image}"
    image_urls.append(image_url)

이제 배치 파일을 만들 수 있어요.

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

모든 게 준비됐으니 API에 업로드해요.

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

파일은 업로드됐지만 아직 배치 추론이 시작되진 않았어요. 시작하려면 job을 만들어야 합니다.

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

배치가 준비되어 실행 중이에요! 정보는 다음 메서드로 조회할 수 있습니다.

retrieved_job = client.batch.jobs.get(job_id=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}%"
)

이 피드백 루프를 자동화해서 결과가 준비되면 다운로드할 수 있도록 해 볼게요.

import time
from IPython.display import clear_output

while retrieved_job.status in ["QUEUED", "RUNNING"]:
    retrieved_job = client.batch.jobs.get(job_id=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)

완료! 이 방법으로 OCR 작업을 굉장히 비용 효율적으로 대량 처리할 수 있어요.

더 알아보기 (Learn more)

• Mistral Vision 문서
• Mistral OCR 쿡북
• HuggingFaceM4/DocumentVQA 데이터셋