Classifying Documents

from typing import Literal
from pydantic import BaseModel, Field

class DocumentClassification(BaseModel):
    rationale: str = Field(..., description="A rationale for the classification, based on the content and visual features of the document. Keep it short and concise, yet detailed enough to justify the classification.")
    document_type: Literal["invoice", "bank-statement", "utility-bill", "other"] = Field(..., description="The type of document being processed")
    confidence: Literal["hi", "med", "lo"] = Field(..., description="Confidence score for the classification, based on the rationale provided and the visual features of the document. For ambiguous documents, the confidence score should be `lo`.")

{
  "rationale": "The document contains a clear 'INVOICE' header, itemized list of products/services, and total amount due. The layout matches typical invoice formats with company details at the top and payment terms at the bottom.",
  "document_type": "invoice",
  "confidence": "hi"
}

# Same imports as before
# ...

# Classify the documents
requests = {}
for path in Path("path/to/documents").glob("*.pdf"):
    prediction: PredictionResponse = client.document.generate(
        file=path,
        domain="document.classification",
        config=GenerationConfig(response_model=DocumentClassification),
        batch=True
    )
    requests[prediction.id] = {
        "path": path,
        "id": prediction.id
    }

# Wait for all predictions to complete
predictions = {}

start_time = time.time()
while time.time() - start_time < 180:
    # fetch the prediction result if it's completed
    for id, request in requests.items():
        # `client.predictions.wait()` will block until the prediction id is completed
        prediction: PredictionResponse = client.predictions.wait(id=id, timeout=10)
        predictions[id] = {**request, "response": prediction.response}

    # wait for 1 second before checking again
    time.sleep(1)

    # break if all predictions are completed
    if len(predictions) == len(requests):
        break

# Get the results
for p in predictions.values():
    print(p)