from retab import Retab

client = Retab()

# Split a multi-page document into sections based on categories
# Use partition_key to identify individual items within a category
result = client.documents.split(
    document="invoice_batch.pdf",
    model="retab-small",
    categories=[
        {"name": "invoice", "description": "Invoice documents with billing information", "partition_key": "invoice number"},
        {"name": "receipt", "description": "Receipt documents for payments"},
        {"name": "contract", "description": "Legal contract documents"},
    ]
)

# Access split results
for split in result.splits:
    print(f"{split.name}: pages {split.pages}")
    # Access partitions when partition_key is specified
    for partition in split.partitions:
        print(f"  {partition.key}: pages {partition.pages}")

# Example output:
# invoice: pages [1, 2, 3]
#   INV-001: pages [1, 2]
#   INV-002: pages [3]
# receipt: pages [4, 5]
# contract: pages [6, 7, 8]

{
  "splits": [
    {
      "name": "invoice",
      "pages": [1, 2, 3],
      "partitions": [
        {
          "key": "INV-001",
          "pages": [1, 2],
          "first_page_y_start": 0.0,
          "last_page_y_end": 1.0
        },
        {
          "key": "INV-002",
          "pages": [3],
          "first_page_y_start": 0.0,
          "last_page_y_end": 1.0
        }
      ]
    },
    {
      "name": "receipt",
      "pages": [4, 5],
      "partitions": []
    },
    {
      "name": "contract",
      "pages": [6, 7, 8],
      "partitions": []
    }
  ]
}