Mistral Batch API로 여러 영수증에서 정보 추출하기
Mistral Batch API로 여러 영수증에서 정보 추출하기
Mistral 배치(batch) API를 사용해 여러 영수증 이미지에서 정보를 추출하고 결과를 pandas DataFrame으로 반환하는 예제를 안내하는 노트북이에요.
출처: 문서
본문
이 노트북에서는 Mistral batch API로 여러 영수증에서 정보를 추출하고, 그 데이터를 pandas DataFrame으로 반환하는 예제를 안내해요. 한 장의 이미지를 Pixtral Large로 처리하는 방법에서 시작해, 배치 잡(batch job)으로 여러 이미지를 한꺼번에 처리하는 흐름까지 따라가요.
!pip install mistralai datasets
영수증 이미지 가져오기
Hugging Face에서 영수증 이미지 데이터셋을 찾았어요:
import pandas as pd
from datasets import load_dataset
# Replace 'dataset_name' with the actual name of the dataset you want to download
dataset_name = 'shirastromer/supermarket-receipts' # Example: IMDB dataset
# Load the dataset
dataset = load_dataset(dataset_name)
# Convert the dataset to a pandas DataFrame
# Assuming you want to load the 'train' split of the dataset
df = pd.DataFrame(dataset['train'])
# Display the first few rows of the DataFrame
df.head()
이미지 하나 살펴보기
한 장의 이미지로 시작해서 Pixtral Large로 그 이미지에서 정보를 얻어볼게요.
# take a look at an image
df.image[1]
Mistral API를 사용해 단일 이미지에서 정보를 추출해요:
import base64
from io import BytesIO
from typing import Any
from PIL.Image import Image
def format_image(image: Image) -> str:
"""
Converts an image to a base64-encoded string with a JPEG format.
Args:
image (Image): The image to be formatted.
Returns:
str: The base64-encoded string with a data URI prefix.
"""
# Convert image to base64
buffer = BytesIO()
image.save(buffer, format="JPEG")
image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
# Add the prefix for base64 format
formatted_base64 = f"data:image/jpeg;base64,{image_base64}"
return formatted_base64
from mistralai.client import Mistral
import os
api_key = os.environ["MISTRAL_API_KEY"]
client = Mistral(api_key=api_key)
# Define the messages for the chat
# Let's extract name, price, and get category for the item
messages = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": format_image(df.image[1])
},
{
"type": "text",
"text": "Extract the name and price of each item on the receipt, categorize each item into one of the following categories: 'Medical', 'Food', 'Beverage', 'Travel', or 'Other', and return the results as a well-structured JSON object. The JSON object should include only the fields: name, price, and classification for each item."
}
]
},
{"role": "assistant", "content": "{", "prefix": True},
]
# Get the chat response
chat_response = client.chat.complete(
model="pixtral-large-latest",
messages=messages,
response_format = {
"type": "json_object",
}
)
# Print the content of the response
print(chat_response.choices[0].message.content)
pixtral-large-latest 모델에 이미지를 base64 인코딩해 보내고, response_format을 json_object로 지정해 구조화된 JSON을 받는 흐름이에요.
Batch API로 여러 이미지 처리하기
배치(batch) 만들기
예시로 10개의 이미지를 처리해 볼게요.
import json
from io import BytesIO
num_samples = 10
list_of_json = []
for idx in range(num_samples):
request = {
"custom_id": str(idx),
"body": {
"max_tokens": 1000,
"messages": [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": format_image(df.image[idx])
},
{
"type": "text",
"text": "Identify the name and price of each item on the receipt, categorize each item into one of the following categories: 'Medical', 'Food', 'Beverage', 'Travel', or 'Other', and return the results as a well-structured JSON object. The JSON object should include only the fields: name, price, and classification for each item."
}
]
},
{
"role": "assistant",
"content": "{",
"prefix": True
}
],
"response_format": {"type": "json_object"}
}
}
list_of_json.append(json.dumps(request).encode("utf-8"))
배치 업로드하기
batch_data = client.files.upload(
file={
"file_name": "file.jsonl",
"content": b"\n".join(list_of_json)},
purpose = "batch"
)
batch_data
각 요청은 custom_id + body(chat 완료 요청)로 구성된 JSONL 라인이고, 이를 파일로 업로드해요.
배치 잡 만들기
created_job = client.batch.jobs.create(
input_files=[batch_data.id],
model="pixtral-large-latest",
endpoint="/v1/chat/completions",
metadata={"job_type": "testing"}
)
created_job
배치 잡 상세 정보 가져오기
retrieved_job = client.batch.jobs.get(job_id=created_job.id)
retrieved_job
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}")
배치 결과 가져오기
output = client.files.download(file_id=retrieved_job.output_file).read().decode("utf-8").strip()
print(output)
정보를 Pandas DataFrame으로 추출하기
# Parse JSON lines
lines = output.strip().split('\n')
# Extract required fields
extracted_data = []
for line in lines:
parsed_line = json.loads(line)
custom_id = parsed_line.get("custom_id")
response = parsed_line.get("response", {})
body = response.get("body", {})
choices = body.get("choices", [])
for choice in choices:
message_content = choice.get("message", {}).get("content", "")
# Extract items from the JSON string in "content"
try:
items_data = json.loads(message_content.strip('```'))
items = items_data if isinstance(items_data, list) else items_data.get("items", [])
for item in items:
extracted_data.append({
"custom_id": custom_id,
"name": item.get("name"),
"price": item.get("price"),
"classification": item.get("classification")
})
except json.JSONDecodeError:
continue
# Create a Pandas DataFrame
df_output = pd.DataFrame(extracted_data)
df_output
각 라인을 파싱해 custom_id, name, price, classification 필드를 모아 DataFrame으로 만든 결과를 확인할 수 있어요.