Split a PDF into documents
A tax return filed as one 40-page PDF holds 19 documents: the federal return, its schedules, supporting statements, and two state returns. Find where each one starts in two steps:
- Ask TypeLLM of every page, in parallel, whether it starts a document.
- Go through the likely starts in order. Show each one beside the first page of the current document, and ask whether it starts a new document or continues the current one.
Set up
Install the library, PyMuPDF to read the PDF, and the Hugging Face client to download it:
pip install "typellm>=0.7.2" pymupdf huggingface_hubGet an API key from your dashboard. The return, its ground truth and the splitting rules come from the PoliTax-Split dataset:
import json
import pymupdf
from huggingface_hub import hf_hub_download
from typellm import TypeLLMClient
client = TypeLLMClient(api_key="YOUR_API_KEY", timeout=90)
REPO = "Extend-AI/PoliTax-Split"
NAME = "harris_2023_federal_state_returns.pdf"
pdf = hf_hub_download(REPO, f"pdfs/harris/{NAME}", repo_type="dataset")
truth = next(row for row in map(json.loads, open(hf_hub_download(REPO, "annotations.jsonl", repo_type="dataset")))
if row["file_name"] == NAME)
taxonomy = json.load(open(hf_hub_download(REPO, "taxonomy.json", repo_type="dataset")))
doc = pymupdf.open(pdf)Write the rules
What counts as one document is up to you. Is a schedule its own document, or part of the return? Are ten numbered statements one document or ten? Write your rules down. Here they come with the dataset: instructions for splitting, and a description of each of its 51 document types.
RULES = ("A PDF holds several documents, one after another: tax forms, schedules and supporting pages. "
"It is split into sub-documents by these rules:\n\n" + taxonomy["splitting_instructions"]
+ "\n\nThe sub-document types:\n\n" + "\n\n".join(f"{t['type']}: {t['description']}" for t in taxonomy["types"]))Ask each page
Send each page's image and text with the rules, and ask one yes/no question with its probability. Every page that scores at least 0.2 is a likely start. The bar is low on purpose: a start missed here can't be recovered later, while an extra one is checked in the next step.
from concurrent.futures import ThreadPoolExecutor
def page_text(i):
text = doc[i].get_text().strip()
return text if len(text) >= 20 else "(no text layer; read the image)"
def image(i):
return doc[i].get_pixmap(dpi=120).tobytes("png")
def yes(context, images, question, thinking=None):
spec = {"type": "boolean", "return_probabilities": True, "instructions": question}
if thinking:
spec["thinking"] = thinking
return client.generate(context=context, images=images, questions={"yes": spec}).result["yes"]["probabilities"][True]
def starts(i):
context = RULES + "\n\nBelow is one page of the PDF: its image is attached, and its text follows.\n\n" + page_text(i)
return yes(context, [image(i)], "Is this page the first page of a sub-document?")
with ThreadPoolExecutor(4) as pool:
probs = list(pool.map(starts, range(doc.page_count)))
probs[0] = 1.0 # the first page always starts one
candidates = [i + 1 for i, p in enumerate(probs) if p >= 0.2]34 likely starts: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 24, 26, 27, 28, 29, 30, 31, 34, 36, 37, 38, 39, 40]All 19 true starts are among the 34, along with 15 extra pages. Some extras score high on their own. Pages 16 to 20 each open with the next numbered statement, often for a different schedule, so each looks like a start at 0.88 to 0.94. But the rules group them with page 15, Statement 1, into one set of supporting statements. Telling them apart takes a second look.
Go through them in order
A page on its own can't always tell whether it starts a new document. A page titled "Statement 2" may continue the statements, and a page with a form's title may be the form's next copy. So show each likely start beside the first page of the current document, and let TypeLLM decide how much to think.
If the likely start begins a new document, it becomes the current document. If it continues the current one, it is dropped, and the next likely start is checked against the same first page:
def new_after(first, later):
context = (RULES + "\n\nBelow are two pages of the PDF, in order: the first page of a sub-document, then a later page. "
"Their images are attached, and their text follows.\n\nThe first page's text:\n" + page_text(first - 1)
+ "\n\nThe later page's text:\n" + page_text(later - 1))
return yes(context, [image(first - 1), image(later - 1)],
"Does the later page start a new sub-document, rather than belong to the sub-document the first page starts?",
thinking="auto")
found = [1]
for page in candidates[1:]:
p = new_after(found[-1], page)
print(f"page {found[-1]} vs {page}: {p:.2f} {'new' if p >= 0.5 else 'same'}")
if p >= 0.5:
found.append(page)page 14 vs 15: 1.00 new
page 15 vs 16: 0.39 same
page 15 vs 17: 0.09 same
page 15 vs 18: 0.29 same
page 15 vs 19: 0.14 same
page 15 vs 20: 0.34 same
page 15 vs 21: 1.00 new
page 21 vs 22: 0.00 same
...
page 37 vs 40: 0.99 newEach step depends on the one before, so these 33 calls run one at a time, at about 30 seconds each with thinking.
Final result
Each document runs from its start to the page before the next start:
ranges = list(zip(found, [s - 1 for s in found[1:]] + [doc.page_count]))Form 1040 pages 1-2: found
Schedule 1 (Form 1040) pages 3-4: found
Schedule 2 (Form 1040) pages 5-6: found
...
supplement pages 15-20: found
Form 540 pages 21-26: found
CA Schedule W-2 pages 27-27: found
CA Schedule W-2 pages 28-28: found
CA Schedule W-2 pages 29-29: found
Schedule CA (540) pages 30-35: found
Schedule S pages 36-36: found
Form D-40 pages 37-39: found
supplement pages 40-40: found
found but not there: []All 19 documents come out with both ends right, and nothing extra. The 73 calls read 499,383 input tokens and think 28,444 tokens: $0.039 for the return, or $0.001 a page at $0.05 per million input tokens and $0.50 per million thinking tokens.
The notebook runs it with your key: open it in Colab.