diff --git a/.env.example b/.env.example index a7ad1cc..414d8f4 100644 --- a/.env.example +++ b/.env.example @@ -1,19 +1,28 @@ -# document-extractor example configuration. Copy to .env and edit. +# document-extractor configuration. Copy to .env and edit. # Placeholder values only. Never commit real keys. -# Choose a backend: "api" (bring your own key) or "hf" (open-source on Hugging Face) -MS_BACKEND=api +# Choose a backend: "api" (any OpenAI-compatible endpoint) or "hf" (open models via transformers) +MS_BACKEND=hf -# --- (a) Bring your own API key: any OpenAI-compatible endpoint --- -OPENAI_BASE_URL=https://your-provider.example/v1 -OPENAI_API_KEY=your-api-key -LLM_MODEL=your-chat-model -EMBEDDING_MODEL=your-embedding-model +# --- (a) api: bring your own key, any OpenAI-compatible endpoint --- +MS_API_BASE=https://your-provider.example/v1 +MS_API_KEY=your-api-key +MS_MODEL=your-chat-model +# Optional: a vision-capable model on the same endpoint, used for OCR of scanned pages +MS_VISION_MODEL= +# Optional: an embedding model on the same endpoint for Q&A retrieval (BM25 is used if empty) +MS_EMBED_MODEL= -# --- (b) Open-source models from Hugging Face --- -HF_MODE=local # local | inference -HF_TOKEN=your-hf-token # only needed for gated models or HF Inference -HF_LLM_MODEL=Qwen/Qwen2.5-7B-Instruct -HF_EMBEDDING_MODEL=BAAI/bge-m3 -HF_RERANKER_MODEL=BAAI/bge-reranker-v2-m3 -HF_OCR_MODEL=microsoft/Florence-2-large +# --- (b) hf: open-source models from Hugging Face, run locally with transformers --- +# Qwen2.5-7B-Instruct needs a GPU or ~32 GB RAM. On a CPU-only machine use the 1.5B model. +MS_HF_MODEL=Qwen/Qwen2.5-7B-Instruct +# MS_HF_MODEL=Qwen/Qwen2.5-1.5B-Instruct +MS_HF_EMBED_MODEL=BAAI/bge-m3 +MS_HF_OCR_MODEL=florence-community/Florence-2-large +# HF_TOKEN=your-hf-token # only needed for gated models + +# --- shared --- +# Q&A retrieval: auto (embeddings when available, else BM25) | embed | bm25 +MS_RETRIEVER=auto +# Where the web demo keeps saved results +MS_CACHE_DIR=.cache/results diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..a78dfee --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,39 @@ +name: CI + +on: + push: + branches: [main, v1] + pull_request: + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12"] + steps: + - uses: actions/checkout@v5 + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install (no model backends; tests use a mocked LLM) + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + - name: Samples are reproducible + run: | + python scripts/make_samples.py + git diff --exit-code -- samples/ + - name: Tests + run: pytest -q + - name: CLI smoke test + run: | + document-extractor --version + # api backend without a vision model: a scanned PDF must fail with a clear message (exit 2) + MS_BACKEND=api MS_API_BASE=http://127.0.0.1:9/v1 MS_MODEL=none \ + document-extractor extract samples/invoice-acme-scanned.pdf --schema invoice || test $? -eq 2 diff --git a/README.md b/README.md index 30e96c2..49fe82a 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@

Turn PDF invoices and contracts into structured fields, each linked to its source.

- Status: in development + Status: v1 License: MIT Python 3.11+ Backends: API or Hugging Face @@ -15,13 +15,18 @@ ## Features -Planned for v1. Progress is tracked in the Roadmap below. +All five features below work in v1. See [Limitations](#limitations) for what they do not cover yet. - **Field extraction to JSON**: define the fields you want (invoice number, totals, parties, dates) and get typed JSON back. - **Source citation per value**: every value carries the page and bounding box it came from, so it can be checked in one click. - **Q&A over the document (RAG)**: ask a question and get an answer grounded in cited passages. -- **Batch mode**: point it at a folder and get one JSON file and one CSV summary per run. -- **Web demo**: drop a PDF, see fields highlighted on the page. +- **Batch mode**: point it at a folder and get one CSV or JSON file per run. +- **Web demo**: pick a sample or upload a PDF, see fields highlighted on the page. + +

+ Web demo: a synthetic invoice is selected, the extracted fields appear with their source boxes on the page, and a question is answered with a cited passage +

+

Web demo on a synthetic invoice (MP4). Results in the recording were computed earlier by Qwen2.5-1.5B-Instruct on CPU and loaded from the demo's result cache.

## Architecture @@ -29,10 +34,11 @@ Planned for v1. Progress is tracked in the Roadmap below. %%{init: {"theme":"base","flowchart":{"curve":"basis","wrappingWidth":260},"themeVariables":{"fontFamily":"Georgia, Times New Roman, serif","fontSize":"14px","primaryColor":"#f3efe6","primaryTextColor":"#1c1a17","primaryBorderColor":"#1c1a17","lineColor":"#1c1a17","clusterBkg":"#faf7f0","clusterBorder":"#8a3a24","titleColor":"#8a3a24","edgeLabelBackground":"#f3efe6"}}}%% flowchart TB P["PDF / scan"]:::ink --> L["Layout and OCR
pages, blocks, boxes"]:::ink + L --> X["Field extractor
schema to JSON"]:::ink L --> C["Chunk and embed"]:::ink - C --> X["Field extractor
schema to JSON"]:::ink C --> Q["Q&A (RAG)
retrieve, answer, cite"]:::ink - X --> O["JSON + source links
page, box, snippet"]:::ink + X --> G["Grounding
match each value to page words"]:::ink + G --> O["JSON + source links
page, box, snippet"]:::ink Q --> O subgraph BK["Backend · pick one"] direction LR @@ -54,8 +60,10 @@ Pick one with `MS_BACKEND` in `.env`. The rest of the pipeline does not change. | | (a) Bring your own API key | (b) Open-source on Hugging Face | |---|---|---| -| **What** | Any OpenAI-compatible endpoint | Open models, run locally or via Hugging Face Inference | -| **Set** | `MS_BACKEND=api`, `OPENAI_BASE_URL`, `OPENAI_API_KEY` | `MS_BACKEND=hf`, `HF_MODE=local` or `inference` | +| **What** | Any OpenAI-compatible endpoint | Open models, run locally with `transformers` | +| **Set** | `MS_BACKEND=api`, `MS_API_BASE`, `MS_API_KEY`, `MS_MODEL` | `MS_BACKEND=hf`, optionally `MS_HF_MODEL` | +| **Scanned pages** | `MS_VISION_MODEL` (a vision model on the same endpoint) | Florence-2 | +| **Q&A retrieval** | `MS_EMBED_MODEL` if set, otherwise BM25 | bge-m3 | | **Good for** | Fastest start, strongest hosted models | Data stays on your machines, no per-call cost | **Default open-source models** @@ -64,31 +72,96 @@ Pick one with `MS_BACKEND` in `.env`. The rest of the pipeline does not change. |---|---|---| | [`Qwen/Qwen2.5-7B-Instruct`](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) | field extraction and answers | Apache-2.0 | | [`BAAI/bge-m3`](https://huggingface.co/BAAI/bge-m3) | embeddings for retrieval | MIT | -| [`BAAI/bge-reranker-v2-m3`](https://huggingface.co/BAAI/bge-reranker-v2-m3) | reranking retrieved passages | Apache-2.0 | -| [`microsoft/Florence-2-large`](https://huggingface.co/microsoft/Florence-2-large) | OCR for scanned pages | MIT | +| [`florence-community/Florence-2-large`](https://huggingface.co/florence-community/Florence-2-large) | OCR for scanned pages (`transformers`-native copy of [`microsoft/Florence-2-large`](https://huggingface.co/microsoft/Florence-2-large)) | MIT | + +On a machine without a GPU, set `MS_HF_MODEL=Qwen/Qwen2.5-1.5B-Instruct`. The 7B model needs a GPU or a lot of memory. +`MS_RETRIEVER=bm25` skips the bge-m3 download and uses keyword retrieval instead. ## Quickstart -> **Coming soon.** The commands below show the planned interface. +Python 3.11 or newer. ```bash git clone https://github.com/ModelShifts/document-extractor cd document-extractor -cp .env.example .env # choose a backend -pip install -e . -document-extractor extract invoice.pdf --schema schemas/invoice.yaml -document-extractor ask contract.pdf "When does the notice period start?" +python -m venv .venv && source .venv/bin/activate +pip install -e ".[hf,demo]" # open models; use ".[api,demo]" for the API backend +cp .env.example .env # choose a backend and models ``` +On a CPU-only machine, install the CPU build of PyTorch first (`pip install torch --index-url https://download.pytorch.org/whl/cpu`) and set `MS_HF_MODEL=Qwen/Qwen2.5-1.5B-Instruct` in `.env`. Models download from Hugging Face on first use. + +```bash +# fields as JSON, each with {page, bbox, text} +document-extractor extract samples/invoice-northwind-classic.pdf --schema schemas/invoice.yaml + +# a question, answered with cited passages +document-extractor ask samples/contract-master-services.pdf "When does the notice period start?" + +# a folder to one CSV (or .json) +document-extractor batch samples --glob "invoice-*.pdf" --schema invoice --out results.csv + +# the web demo at http://127.0.0.1:7860 (python app.py does the same) +document-extractor serve +``` + +### Using the API backend + +Any OpenAI-compatible endpoint works. In `.env`: + +```bash +MS_BACKEND=api +MS_API_BASE=https://your-provider.example/v1 +MS_API_KEY=your-api-key +MS_MODEL=your-chat-model +``` + +To try it without a paid key, run an open model behind a local OpenAI-compatible server, for example [llama-cpp-python](https://github.com/abetlen/llama-cpp-python): + +```bash +pip install "llama-cpp-python[server]" huggingface_hub \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu # prebuilt CPU wheels, no compiler needed +hf download Qwen/Qwen2.5-1.5B-Instruct-GGUF qwen2.5-1.5b-instruct-q4_k_m.gguf --local-dir models +python -m llama_cpp.server --model models/qwen2.5-1.5b-instruct-q4_k_m.gguf --n_ctx 8192 --port 8001 +# then: MS_API_BASE=http://127.0.0.1:8001/v1 MS_API_KEY=local MS_MODEL=qwen2.5-1.5b-instruct-q4_k_m +``` + +### Schemas + +A schema is a YAML file that lists fields and their types (`string`, `date`, `number`, `integer`). See [`schemas/invoice.yaml`](schemas/invoice.yaml) and [`schemas/contract.yaml`](schemas/contract.yaml). Dates come back as `YYYY-MM-DD` and amounts as numbers. The model is asked to copy each value as printed, and the value is then matched back to the words on the page to find its box. + +### Sample documents and tests + +`samples/` holds seven synthetic PDFs (four invoice layouts, two contracts, one image-only scan) with a ground-truth file for each. All names and figures in them are fictional. + +```bash +pip install -e ".[dev]" +python scripts/make_samples.py # regenerate samples/ +pytest # uses a mocked model, no downloads +python scripts/evaluate.py # compare the configured backend with the ground truth +``` + +## Limitations + +- The samples are synthetic and have simple layouts. The project has not been measured on real-world documents. +- On a CPU, extraction with the open models takes tens of seconds per document, and loading Qwen, bge-m3 and Florence-2 together needs several GB of memory. The 7B default needs a GPU. +- Florence-2 reads a scanned page in one pass, so small or dense text can be misread. On the API backend, OCR through `MS_VISION_MODEL` is implemented but has only been tested with a mocked model. +- Source boxes come from matching the value against the page text. If a value is printed more than once, the copy closest to a matching label is used, which is not always the one the model read. +- Numeric dates such as `03.02.2026` are read day-first. +- The web demo saves results under `.cache/results` and reuses them for the same file, schema and model. + ## Roadmap - [x] Repository scaffold and README -- [ ] PDF parsing with page and box coordinates -- [ ] Schema-driven field extraction -- [ ] Source citation per value -- [ ] Q&A over the document (RAG) -- [ ] Batch mode -- [ ] Web demo +- [x] PDF parsing with page and box coordinates +- [x] Schema-driven field extraction +- [x] Source citation per value +- [x] Q&A over the document (RAG) +- [x] Batch mode +- [x] Web demo +- [ ] Reranking retrieved passages with `BAAI/bge-reranker-v2-m3` +- [ ] Scanned-page OCR tested with a real vision model on the API backend +- [ ] Hosted demo ## License diff --git a/app.py b/app.py new file mode 100644 index 0000000..bf7f6d5 --- /dev/null +++ b/app.py @@ -0,0 +1,11 @@ +"""Entry point for the web demo, also usable as a Hugging Face Space app file. + + python app.py # same as: document-extractor serve +""" + +import os + +from document_extractor.webapp import launch + +if __name__ == "__main__": + launch(host=os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1"), port=int(os.environ.get("GRADIO_SERVER_PORT", 7860))) diff --git a/assets/banner.svg b/assets/banner.svg index 067715b..e607037 100644 --- a/assets/banner.svg +++ b/assets/banner.svg @@ -36,7 +36,7 @@ MODELSHIFTS 01 / 06 - IN DEVELOPMENT + V1 MIT diff --git a/assets/demo.gif b/assets/demo.gif new file mode 100644 index 0000000..d322e89 Binary files /dev/null and b/assets/demo.gif differ diff --git a/assets/demo.mp4 b/assets/demo.mp4 new file mode 100644 index 0000000..2dac658 Binary files /dev/null and b/assets/demo.mp4 differ diff --git a/pyproject.toml b/pyproject.toml index e6970f5..6ce2b26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,17 +4,29 @@ build-backend = "hatchling.build" [project] name = "document-extractor" -version = "0.0.1" +version = "0.1.0" description = "Turn PDF invoices and contracts into structured fields, each linked to its source." readme = "README.md" license = "MIT" requires-python = ">=3.11" authors = [{ name = "Modelshifts", email = "info@modelshifts.com" }] -dependencies = ["pydantic>=2", "httpx>=0.27"] +dependencies = [ + "pydantic>=2", + "pdfplumber>=0.11", + "pypdfium2>=4.30", + "pillow>=10", + "pyyaml>=6", + "rank-bm25>=0.2.2", + "python-dotenv>=1", + "numpy>=1.26", +] [project.optional-dependencies] api = ["openai>=1.40"] -hf = ["transformers>=4.45", "sentence-transformers>=3", "torch>=2.3"] +hf = ["transformers>=4.56", "sentence-transformers>=3", "torch>=2.3", "accelerate>=0.33"] +demo = ["gradio>=5"] +samples = ["reportlab>=4"] +dev = ["pytest>=8", "reportlab>=4", "openai>=1.40"] [project.urls] Homepage = "https://modelshifts.com/work" @@ -25,3 +37,9 @@ document-extractor = "document_extractor.cli:main" [tool.hatch.build.targets.wheel] packages = ["src/document_extractor"] + +[tool.hatch.build.targets.wheel.force-include] +"schemas" = "document_extractor/schemas" + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/samples/contract-master-services.pdf b/samples/contract-master-services.pdf new file mode 100644 index 0000000..b8fc918 Binary files /dev/null and b/samples/contract-master-services.pdf differ diff --git a/samples/contract-master-services.truth.json b/samples/contract-master-services.truth.json new file mode 100644 index 0000000..770896f --- /dev/null +++ b/samples/contract-master-services.truth.json @@ -0,0 +1,14 @@ +{ + "file": "contract-master-services.pdf", + "schema": "contract", + "fields": { + "contract_title": "Master Services Agreement", + "party_a": "Northwind Parts Pte. Ltd.", + "party_b": "Example Logistics Pte. Ltd.", + "effective_date": "2026-07-01", + "term_months": 24, + "notice_period_days": 60, + "payment_terms_days": 30, + "governing_law": "Singapore" + } +} diff --git a/samples/contract-supply.pdf b/samples/contract-supply.pdf new file mode 100644 index 0000000..51d3200 Binary files /dev/null and b/samples/contract-supply.pdf differ diff --git a/samples/contract-supply.truth.json b/samples/contract-supply.truth.json new file mode 100644 index 0000000..27faabb --- /dev/null +++ b/samples/contract-supply.truth.json @@ -0,0 +1,14 @@ +{ + "file": "contract-supply.pdf", + "schema": "contract", + "fields": { + "contract_title": "Supply Agreement", + "party_a": "Sample Robotics Ltd.", + "party_b": "Demo Freight GmbH", + "effective_date": "2026-09-15", + "term_months": 12, + "notice_period_days": 90, + "payment_terms_days": 45, + "governing_law": "England and Wales" + } +} diff --git a/samples/invoice-acme-scanned.pdf b/samples/invoice-acme-scanned.pdf new file mode 100644 index 0000000..cf654f0 Binary files /dev/null and b/samples/invoice-acme-scanned.pdf differ diff --git a/samples/invoice-acme-scanned.truth.json b/samples/invoice-acme-scanned.truth.json new file mode 100644 index 0000000..beb7294 --- /dev/null +++ b/samples/invoice-acme-scanned.truth.json @@ -0,0 +1,15 @@ +{ + "file": "invoice-acme-scanned.pdf", + "schema": "invoice", + "fields": { + "invoice_number": "AES-5120", + "invoice_date": "2026-06-08", + "due_date": "2026-07-08", + "seller_name": "Acme Example Services Pte. Ltd.", + "buyer_name": "Northwind Parts Pte. Ltd.", + "currency": "SGD", + "subtotal": 5219.37, + "tax": 469.74, + "total": 5689.11 + } +} diff --git a/samples/invoice-demo-freight-compact.pdf b/samples/invoice-demo-freight-compact.pdf new file mode 100644 index 0000000..258d6b4 Binary files /dev/null and b/samples/invoice-demo-freight-compact.pdf differ diff --git a/samples/invoice-demo-freight-compact.truth.json b/samples/invoice-demo-freight-compact.truth.json new file mode 100644 index 0000000..43e8263 --- /dev/null +++ b/samples/invoice-demo-freight-compact.truth.json @@ -0,0 +1,15 @@ +{ + "file": "invoice-demo-freight-compact.pdf", + "schema": "invoice", + "fields": { + "invoice_number": "DF-88213", + "invoice_date": "2026-02-03", + "due_date": "2026-03-05", + "seller_name": "Demo Freight GmbH", + "buyer_name": "Northwind Parts Pte. Ltd.", + "currency": "EUR", + "subtotal": 3151.06, + "tax": 598.7, + "total": 3749.76 + } +} diff --git a/samples/invoice-example-logistics-banner.pdf b/samples/invoice-example-logistics-banner.pdf new file mode 100644 index 0000000..922a65e Binary files /dev/null and b/samples/invoice-example-logistics-banner.pdf differ diff --git a/samples/invoice-example-logistics-banner.truth.json b/samples/invoice-example-logistics-banner.truth.json new file mode 100644 index 0000000..7e35558 --- /dev/null +++ b/samples/invoice-example-logistics-banner.truth.json @@ -0,0 +1,15 @@ +{ + "file": "invoice-example-logistics-banner.pdf", + "schema": "invoice", + "fields": { + "invoice_number": "EL/INV/7731", + "invoice_date": "2026-05-02", + "due_date": "2026-06-01", + "seller_name": "Example Logistics Inc.", + "buyer_name": "Sample Robotics Ltd.", + "currency": "USD", + "subtotal": 12594.58, + "tax": 881.62, + "total": 13476.2 + } +} diff --git a/samples/invoice-northwind-classic.pdf b/samples/invoice-northwind-classic.pdf new file mode 100644 index 0000000..c883752 Binary files /dev/null and b/samples/invoice-northwind-classic.pdf differ diff --git a/samples/invoice-northwind-classic.truth.json b/samples/invoice-northwind-classic.truth.json new file mode 100644 index 0000000..32011b2 --- /dev/null +++ b/samples/invoice-northwind-classic.truth.json @@ -0,0 +1,15 @@ +{ + "file": "invoice-northwind-classic.pdf", + "schema": "invoice", + "fields": { + "invoice_number": "NWP-2026-0412", + "invoice_date": "2026-03-14", + "due_date": "2026-04-13", + "seller_name": "Northwind Parts Pte. Ltd.", + "buyer_name": "Example Logistics Pte. Ltd.", + "currency": "SGD", + "subtotal": 5042.83, + "tax": 453.85, + "total": 5496.68 + } +} diff --git a/samples/invoice-placeholder-two-page.pdf b/samples/invoice-placeholder-two-page.pdf new file mode 100644 index 0000000..ea20833 Binary files /dev/null and b/samples/invoice-placeholder-two-page.pdf differ diff --git a/samples/invoice-placeholder-two-page.truth.json b/samples/invoice-placeholder-two-page.truth.json new file mode 100644 index 0000000..ccf0e41 --- /dev/null +++ b/samples/invoice-placeholder-two-page.truth.json @@ -0,0 +1,15 @@ +{ + "file": "invoice-placeholder-two-page.pdf", + "schema": "invoice", + "fields": { + "invoice_number": "PS-INV-30027", + "invoice_date": "2026-01-20", + "due_date": "2026-02-19", + "seller_name": "Placeholder Supplies Pte. Ltd.", + "buyer_name": "Sample Robotics Ltd.", + "currency": "SGD", + "subtotal": 6204.22, + "tax": 558.38, + "total": 6762.6 + } +} diff --git a/schemas/contract.yaml b/schemas/contract.yaml new file mode 100644 index 0000000..1db1f1f --- /dev/null +++ b/schemas/contract.yaml @@ -0,0 +1,27 @@ +name: contract +description: A commercial agreement between two companies. +fields: + contract_title: + type: string + description: The title of the agreement as printed at the top + party_a: + type: string + description: The first company named as a party to the agreement + party_b: + type: string + description: The second company named as a party to the agreement + effective_date: + type: date + description: The date the agreement is made or takes effect + term_months: + type: integer + description: Length of the initial term in months + notice_period_days: + type: integer + description: Days of written notice needed to terminate the agreement + payment_terms_days: + type: integer + description: Number of days within which invoices must be paid + governing_law: + type: string + description: The jurisdiction whose laws govern the agreement diff --git a/schemas/invoice.yaml b/schemas/invoice.yaml new file mode 100644 index 0000000..71ff5e6 --- /dev/null +++ b/schemas/invoice.yaml @@ -0,0 +1,30 @@ +name: invoice +description: A commercial invoice from a seller to a buyer. +fields: + invoice_number: + type: string + description: The invoice number or reference printed on the invoice + invoice_date: + type: date + description: The date the invoice was issued + due_date: + type: date + description: The date payment is due + seller_name: + type: string + description: The company that issued the invoice (the supplier or vendor) + buyer_name: + type: string + description: The company the invoice is billed to (the customer) + currency: + type: string + description: Three-letter currency code, for example SGD, USD or EUR + subtotal: + type: number + description: Total before tax (net amount, sub-total) + tax: + type: number + description: Tax amount (GST, VAT or sales tax) + total: + type: number + description: Final amount due including tax (total due, gross total, amount due) diff --git a/scripts/evaluate.py b/scripts/evaluate.py new file mode 100644 index 0000000..c64078a --- /dev/null +++ b/scripts/evaluate.py @@ -0,0 +1,96 @@ +"""Compare extraction results with the ground truth in ``samples/*.truth.json``. + + python scripts/evaluate.py # run the configured backend on every sample + python scripts/evaluate.py --out eval.json # also save the raw results + +A field counts as correct when the typed value matches the truth (strings are +compared case-insensitively with spaces and trailing punctuation ignored; +numbers within 0.005). "Sourced" counts correct values that also carry a source +box whose text matches the value. +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +import time +from pathlib import Path + +from document_extractor.config import Settings +from document_extractor.grounding import _key, variants +from document_extractor.pipeline import run_extract +from document_extractor.schema import load_schema + +ROOT = Path(__file__).resolve().parent.parent + + +def same(pred, truth, ftype: str) -> bool: + if pred is None or truth is None: + return pred is truth + if ftype == "number": + return abs(float(pred) - float(truth)) < 0.005 + if ftype == "integer": + return int(pred) == int(truth) + clean = lambda s: re.sub(r"\s+", " ", str(s)).strip(" .,;:").lower() # noqa: E731 + return clean(pred) == clean(truth) + + +def source_matches(fv, ftype: str) -> bool: + """True when the source box text spells the value (or contains it).""" + if fv.source is None: + return False + got = _key(fv.source.text.split()) + return any(_key(v.split()) and _key(v.split()) in got for v in variants(fv.raw, fv.value, ftype)) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--samples", default=str(ROOT / "samples")) + ap.add_argument("--out") + args = ap.parse_args() + settings = Settings.from_env() + settings.validate() + rows, raw = [], [] + tot_ok = tot_n = tot_src = 0 + for truth_file in sorted(Path(args.samples).glob("*.truth.json")): + truth = json.loads(truth_file.read_text()) + pdf = truth_file.with_name(truth["file"]) + schema = load_schema(truth["schema"]) + t0 = time.time() + try: + result = run_extract(pdf, truth["schema"], settings) + except Exception as exc: # report and keep going + rows.append((pdf.name, 0, len(schema.fields), 0, 0.0, f"error: {exc}")) + tot_n += len(schema.fields) + continue + dt = time.time() - t0 + raw.append(result.to_json_dict()) + ok = src = 0 + misses = [] + for spec in schema.fields: + fv = result.fields[spec.name] + if same(fv.value, truth["fields"][spec.name], spec.type): + ok += 1 + if source_matches(fv, spec.type): + src += 1 + else: + misses.append(f"{spec.name}={fv.value!r} (want {truth['fields'][spec.name]!r})") + rows.append((pdf.name, ok, len(schema.fields), src, dt, "; ".join(misses))) + tot_ok, tot_n, tot_src = tot_ok + ok, tot_n + len(schema.fields), tot_src + src + print(f"{pdf.name}: {ok}/{len(schema.fields)} correct, {src} sourced, {dt:.1f}s", file=sys.stderr) + + print(f"\nBackend: {settings.backend}, model: {settings.model_name}\n") + print("| Sample | Correct | Sourced | Time (s) | Misses |") + print("|---|---|---|---|---|") + for name, ok, n, src, dt, miss in rows: + print(f"| {name} | {ok}/{n} | {src}/{ok} | {dt:.1f} | {miss or '-'} |") + print(f"| **All** | **{tot_ok}/{tot_n} ({100 * tot_ok / max(tot_n, 1):.0f}%)** | {tot_src}/{tot_ok} | | |") + if args.out: + Path(args.out).write_text(json.dumps(raw, indent=2) + "\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/make_samples.py b/scripts/make_samples.py new file mode 100644 index 0000000..cc1be1e --- /dev/null +++ b/scripts/make_samples.py @@ -0,0 +1,589 @@ +"""Generate the synthetic sample documents in ``samples/``. + +Every company, address and number in these files is invented. Nothing here is +real data. Run from the repository root: + + python scripts/make_samples.py + +Each PDF gets a ``.truth.json`` file next to it with the expected field +values, used by the tests and the accuracy check in ``scripts/evaluate.py``. +""" + +from __future__ import annotations + +import io +import json +import random +from pathlib import Path + +from reportlab.lib.colors import HexColor, black, white +from reportlab.lib.pagesizes import A4, letter +from reportlab.lib.utils import ImageReader +from reportlab.pdfgen import canvas + +OUT = Path(__file__).resolve().parent.parent / "samples" +INK = HexColor("#1c1a17") +RUST = HexColor("#8a3a24") +PAPER = HexColor("#f3efe6") +GREY = HexColor("#6b665e") + +rng = random.Random(20260311) + + +def money(x: float) -> str: + return f"{x:,.2f}" + + +def make_items(n: int, catalogue: list[str], lo: float, hi: float) -> list[tuple[str, int, float, float]]: + items = [] + for desc in rng.sample(catalogue, n): + qty = rng.randint(1, 12) + unit = round(rng.uniform(lo, hi), 2) + items.append((desc, qty, unit, round(qty * unit, 2))) + return items + + +def totals(items, rate: float) -> tuple[float, float, float]: + sub = round(sum(i[3] for i in items), 2) + tax = round(sub * rate, 2) + return sub, tax, round(sub + tax, 2) + + +def write_truth(name: str, schema: str, fields: dict) -> None: + (OUT / f"{name}.truth.json").write_text( + json.dumps({"file": f"{name}.pdf", "schema": schema, "fields": fields}, indent=2) + "\n" + ) + + +PARTS = [ + "Hex bolt M8 x 40, zinc", + "Bearing 6204-2RS", + "Timing belt HTD 5M-450", + "Aluminium extrusion 20x20, 1 m", + "Shaft coupling 8 mm to 10 mm", + "Linear rail MGN12, 400 mm", + "Stepper driver board", + "Cable chain 10x15, 1 m", + "Proximity sensor, inductive", + "Spring washer M8, pack of 100", +] +SERVICES = [ + "Pallet storage, per week", + "Last-mile delivery, zone B", + "Customs documentation", + "Cross-dock handling", + "Shrink wrapping, per pallet", + "Forklift hire, per hour", + "Container unloading", + "Labelling service", +] + + +# --------------------------------------------------------------------------- # +# Invoice 1: classic layout, seller top-left, totals bottom-right +# --------------------------------------------------------------------------- # +def invoice_classic(path: Path) -> dict: + c = canvas.Canvas(str(path), pagesize=A4, invariant=1) + w, h = A4 + items = make_items(5, PARTS, 2.5, 180) + sub, tax, total = totals(items, 0.09) + + c.setFillColor(INK) + c.setFont("Helvetica-Bold", 18) + c.drawString(50, h - 70, "Northwind Parts Pte. Ltd.") + c.setFont("Helvetica", 9) + c.drawString(50, h - 86, "12 Example Road, #04-01, Singapore 000000") + c.drawString(50, h - 98, "GST Reg. No. 000000000X (sample)") + + c.setFont("Helvetica-Bold", 24) + c.setFillColor(RUST) + c.drawRightString(w - 50, h - 70, "TAX INVOICE") + c.setFillColor(INK) + c.setFont("Helvetica", 10) + c.drawRightString(w - 50, h - 92, "Invoice No: NWP-2026-0412") + c.drawRightString(w - 50, h - 106, "Invoice Date: 14 March 2026") + c.drawRightString(w - 50, h - 120, "Due Date: 13 April 2026") + + c.setFont("Helvetica-Bold", 10) + c.drawString(50, h - 150, "Bill To") + c.setFont("Helvetica", 10) + c.drawString(50, h - 164, "Example Logistics Pte. Ltd.") + c.drawString(50, h - 177, "88 Sample Avenue, Singapore 000001") + + y = h - 220 + c.setFillColor(PAPER) + c.rect(45, y - 6, w - 90, 20, stroke=0, fill=1) + c.setFillColor(INK) + c.setFont("Helvetica-Bold", 9) + c.drawString(50, y, "Description") + c.drawRightString(360, y, "Qty") + c.drawRightString(450, y, "Unit price") + c.drawRightString(w - 50, y, "Amount") + c.setFont("Helvetica", 9) + for desc, qty, unit, amt in items: + y -= 20 + c.drawString(50, y, desc) + c.drawRightString(360, y, str(qty)) + c.drawRightString(450, y, money(unit)) + c.drawRightString(w - 50, y, money(amt)) + y -= 14 + c.line(300, y, w - 50, y) + c.setFont("Helvetica", 10) + for label, val in (("Subtotal", sub), ("GST 9%", tax)): + y -= 18 + c.drawString(330, y, label) + c.drawRightString(w - 50, y, f"SGD {money(val)}") + y -= 22 + c.setFont("Helvetica-Bold", 11) + c.drawString(330, y, "Total Due") + c.drawRightString(w - 50, y, f"SGD {money(total)}") + + c.setFont("Helvetica", 8) + c.setFillColor(GREY) + c.drawString(50, 60, "Synthetic sample document generated for testing. All names and figures are fictional.") + c.save() + return { + "invoice_number": "NWP-2026-0412", + "invoice_date": "2026-03-14", + "due_date": "2026-04-13", + "seller_name": "Northwind Parts Pte. Ltd.", + "buyer_name": "Example Logistics Pte. Ltd.", + "currency": "SGD", + "subtotal": sub, + "tax": tax, + "total": total, + } + + +# --------------------------------------------------------------------------- # +# Invoice 2: banner header, two-column parties, different labels +# --------------------------------------------------------------------------- # +def invoice_banner(path: Path) -> dict: + c = canvas.Canvas(str(path), pagesize=letter, invariant=1) + w, h = letter + items = make_items(4, SERVICES, 40, 900) + sub, tax, total = totals(items, 0.07) + + c.setFillColor(INK) + c.rect(0, h - 90, w, 90, stroke=0, fill=1) + c.setFillColor(white) + c.setFont("Helvetica-Bold", 22) + c.drawString(40, h - 55, "Example Logistics") + c.setFont("Helvetica", 10) + c.drawRightString(w - 40, h - 45, "INVOICE") + c.drawRightString(w - 40, h - 60, "Bill No. EL/INV/7731") + + c.setFillColor(INK) + c.setFont("Helvetica-Bold", 9) + c.drawString(40, h - 120, "FROM") + c.drawString(w / 2, h - 120, "TO") + c.setFont("Helvetica", 10) + c.drawString(40, h - 135, "Example Logistics Inc.") + c.drawString(40, h - 148, "400 Placeholder Street, Springfield 00000") + c.drawString(w / 2, h - 135, "Sample Robotics Ltd.") + c.drawString(w / 2, h - 148, "7 Demo Park, Unit 3, Testville 00000") + + c.setFont("Helvetica", 10) + c.drawString(40, h - 180, "Issued: 2026-05-02") + c.drawString(200, h - 180, "Payment due: 2026-06-01") + c.drawString(380, h - 180, "Currency: USD") + + y = h - 220 + c.setStrokeColor(INK) + c.line(40, y + 14, w - 40, y + 14) + c.setFont("Helvetica-Bold", 9) + c.drawString(40, y, "Service") + c.drawRightString(400, y, "Units") + c.drawRightString(480, y, "Rate") + c.drawRightString(w - 40, y, "Line total") + c.line(40, y - 6, w - 40, y - 6) + c.setFont("Helvetica", 9) + for desc, qty, unit, amt in items: + y -= 20 + c.drawString(40, y, desc) + c.drawRightString(400, y, str(qty)) + c.drawRightString(480, y, money(unit)) + c.drawRightString(w - 40, y, money(amt)) + + y -= 40 + c.setFillColor(PAPER) + c.rect(330, y - 60, w - 370, 80, stroke=0, fill=1) + c.setFillColor(INK) + c.setFont("Helvetica", 10) + c.drawString(340, y, "Net amount") + c.drawRightString(w - 50, y, money(sub)) + c.drawString(340, y - 18, "Sales tax (7%)") + c.drawRightString(w - 50, y - 18, money(tax)) + c.setFont("Helvetica-Bold", 12) + c.drawString(340, y - 44, "Amount Due (USD)") + c.drawRightString(w - 50, y - 44, money(total)) + + c.setFont("Helvetica", 8) + c.setFillColor(GREY) + c.drawString(40, 50, "Synthetic sample document generated for testing. All names and figures are fictional.") + c.save() + return { + "invoice_number": "EL/INV/7731", + "invoice_date": "2026-05-02", + "due_date": "2026-06-01", + "seller_name": "Example Logistics Inc.", + "buyer_name": "Sample Robotics Ltd.", + "currency": "USD", + "subtotal": sub, + "tax": tax, + "total": total, + } + + +# --------------------------------------------------------------------------- # +# Invoice 3: compact European style, key-value block on the right, EUR +# --------------------------------------------------------------------------- # +def invoice_compact(path: Path) -> dict: + c = canvas.Canvas(str(path), pagesize=A4, invariant=1) + w, h = A4 + items = make_items(3, PARTS, 10, 400) + sub, tax, total = totals(items, 0.19) + + c.setFillColor(RUST) + c.rect(0, 0, 14, h, stroke=0, fill=1) + c.setFillColor(INK) + c.setFont("Times-Bold", 16) + c.drawString(50, h - 60, "Demo Freight GmbH") + c.setFont("Times-Roman", 9) + c.drawString(50, h - 74, "Beispielweg 1, 00000 Musterstadt") + + c.setFont("Times-Roman", 10) + c.drawString(50, h - 130, "Customer:") + c.drawString(50, h - 144, "Northwind Parts Pte. Ltd.") + c.drawString(50, h - 157, "12 Example Road, Singapore 000000") + + kv = [ + ("Invoice number", "DF-88213"), + ("Invoice date", "03.02.2026"), + ("Payable by", "05.03.2026"), + ("Currency", "EUR"), + ] + y = h - 130 + for k, v in kv: + c.setFont("Times-Bold", 10) + c.drawString(340, y, k) + c.setFont("Times-Roman", 10) + c.drawRightString(w - 50, y, v) + y -= 15 + + y = h - 220 + c.setFont("Times-Bold", 10) + c.drawString(50, y, "Pos.") + c.drawString(85, y, "Article") + c.drawRightString(380, y, "Qty") + c.drawRightString(460, y, "Price") + c.drawRightString(w - 50, y, "Total") + c.setFont("Times-Roman", 10) + for i, (desc, qty, unit, amt) in enumerate(items, 1): + y -= 18 + c.drawString(50, y, str(i)) + c.drawString(85, y, desc) + c.drawRightString(380, y, str(qty)) + c.drawRightString(460, y, money(unit)) + c.drawRightString(w - 50, y, money(amt)) + y -= 30 + c.drawString(360, y, "Net total") + c.drawRightString(w - 50, y, f"{money(sub)} EUR") + y -= 15 + c.drawString(360, y, "VAT 19%") + c.drawRightString(w - 50, y, f"{money(tax)} EUR") + y -= 18 + c.setFont("Times-Bold", 11) + c.drawString(360, y, "Gross total") + c.drawRightString(w - 50, y, f"{money(total)} EUR") + + c.setFont("Times-Roman", 8) + c.setFillColor(GREY) + c.drawString(50, 50, "Synthetic sample document generated for testing. All names and figures are fictional.") + c.save() + return { + "invoice_number": "DF-88213", + "invoice_date": "2026-02-03", + "due_date": "2026-03-05", + "seller_name": "Demo Freight GmbH", + "buyer_name": "Northwind Parts Pte. Ltd.", + "currency": "EUR", + "subtotal": sub, + "tax": tax, + "total": total, + } + + +# --------------------------------------------------------------------------- # +# Invoice 4: two pages, long item list, totals on page 2 +# --------------------------------------------------------------------------- # +def invoice_two_page(path: Path) -> dict: + c = canvas.Canvas(str(path), pagesize=A4, invariant=1) + w, h = A4 + catalogue = [f"{p} (batch {k})" for p in PARTS for k in ("A", "B", "C")] + items = make_items(30, catalogue, 1.2, 60) + sub, tax, total = totals(items, 0.09) + + def header(page: int) -> float: + c.setFillColor(INK) + c.setFont("Helvetica-Bold", 14) + c.drawString(50, h - 60, "Placeholder Supplies Pte. Ltd.") + c.setFont("Helvetica", 9) + c.drawRightString(w - 50, h - 60, f"Page {page} of 2") + c.drawRightString(w - 50, h - 72, "Ref: PS-INV-30027") + return h - 100 + + y = header(1) + c.setFont("Helvetica", 10) + c.drawString(50, y, "Invoice to: Sample Robotics Ltd.") + c.drawString(50, y - 14, "Date of issue: 20/01/2026") + c.drawString(50, y - 28, "Terms: Net 30, due 19/02/2026") + y -= 60 + c.setFont("Helvetica-Bold", 9) + c.drawString(50, y, "Item") + c.drawRightString(380, y, "Qty") + c.drawRightString(460, y, "Unit (SGD)") + c.drawRightString(w - 50, y, "Amount (SGD)") + c.setFont("Helvetica", 9) + for n, (desc, qty, unit, amt) in enumerate(items): + if n == 22: + c.showPage() + y = header(2) - 10 + c.setFont("Helvetica", 9) + y -= 17 + c.drawString(50, y, desc) + c.drawRightString(380, y, str(qty)) + c.drawRightString(460, y, money(unit)) + c.drawRightString(w - 50, y, money(amt)) + y -= 30 + c.setFont("Helvetica", 10) + c.drawString(340, y, "Sub-total") + c.drawRightString(w - 50, y, money(sub)) + c.drawString(340, y - 16, "GST @ 9%") + c.drawRightString(w - 50, y - 16, money(tax)) + c.setFont("Helvetica-Bold", 11) + c.drawString(340, y - 36, "TOTAL SGD") + c.drawRightString(w - 50, y - 36, money(total)) + c.setFont("Helvetica", 8) + c.setFillColor(GREY) + c.drawString(50, 50, "Synthetic sample document generated for testing. All names and figures are fictional.") + c.save() + return { + "invoice_number": "PS-INV-30027", + "invoice_date": "2026-01-20", + "due_date": "2026-02-19", + "seller_name": "Placeholder Supplies Pte. Ltd.", + "buyer_name": "Sample Robotics Ltd.", + "currency": "SGD", + "subtotal": sub, + "tax": tax, + "total": total, + } + + +# --------------------------------------------------------------------------- # +# Contracts +# --------------------------------------------------------------------------- # +def paragraph(c: canvas.Canvas, text: str, x: float, y: float, width: float, font="Times-Roman", size=10.5, lead=14) -> float: + c.setFont(font, size) + words, line = text.split(), "" + for word in words: + trial = f"{line} {word}".strip() + if c.stringWidth(trial, font, size) > width: + c.drawString(x, y, line) + y -= lead + line = word + else: + line = trial + if line: + c.drawString(x, y, line) + y -= lead + return y - 6 + + +def contract(path: Path, title: str, intro: str, clauses: list[tuple[str, str]]) -> None: + c = canvas.Canvas(str(path), pagesize=A4, invariant=1) + w, h = A4 + c.setFillColor(INK) + c.setFont("Times-Bold", 16) + c.drawCentredString(w / 2, h - 70, title) + y = paragraph(c, intro, 60, h - 110, w - 120) + for heading, body in clauses: + if y < 140: + c.showPage() + c.setFillColor(INK) + y = h - 70 + c.setFont("Times-Bold", 11) + c.drawString(60, y, heading) + y = paragraph(c, body, 60, y - 16, w - 120) + y -= 20 + c.setFont("Times-Roman", 10) + c.drawString(60, y, "Signed for and on behalf of each party by an authorised signatory.") + c.setFont("Times-Roman", 8) + c.setFillColor(GREY) + c.drawString(60, 50, "Synthetic sample document generated for testing. All names and terms are fictional.") + c.save() + + +def contract_services(path: Path) -> dict: + contract( + path, + "MASTER SERVICES AGREEMENT", + "This Master Services Agreement is made on 1 July 2026 (the Effective Date) between " + "Northwind Parts Pte. Ltd. (the Customer) and Example Logistics Pte. Ltd. (the Provider).", + [ + ("1. Services", "The Provider will supply warehousing and delivery services as described in each " + "statement of work agreed in writing by both parties."), + ("2. Term", "This Agreement starts on the Effective Date and continues for a period of 24 months, " + "unless terminated earlier under clause 6."), + ("3. Fees and Payment", "The Customer will pay each undisputed invoice within 30 days of the invoice " + "date. Fees are stated in Singapore dollars and exclude goods and services tax."), + ("4. Service Levels", "The Provider will use reasonable efforts to meet the delivery windows set out " + "in the applicable statement of work."), + ("5. Confidentiality", "Each party will keep the other party's confidential information private and " + "use it only to perform this Agreement."), + ("6. Termination", "Either party may terminate this Agreement for convenience by giving not less than " + "sixty (60) days written notice to the other party. The notice period starts on the day the notice " + "is received."), + ("7. Governing Law", "This Agreement is governed by the laws of Singapore."), + ], + ) + return { + "contract_title": "Master Services Agreement", + "party_a": "Northwind Parts Pte. Ltd.", + "party_b": "Example Logistics Pte. Ltd.", + "effective_date": "2026-07-01", + "term_months": 24, + "notice_period_days": 60, + "payment_terms_days": 30, + "governing_law": "Singapore", + } + + +def contract_supply(path: Path) -> dict: + contract( + path, + "SUPPLY AGREEMENT", + "This Supply Agreement is entered into as of 15 September 2026 by and between " + "Sample Robotics Ltd. (the Buyer) and Demo Freight GmbH (the Supplier).", + [ + ("1. Supply of Goods", "The Supplier agrees to supply the components listed in Schedule 1 against " + "purchase orders issued by the Buyer."), + ("2. Duration", "This Agreement shall remain in force for twelve (12) months from the date above and " + "renews automatically for further periods of twelve months."), + ("3. Price and Payment", "Prices are fixed for the duration of the Agreement. Invoices are payable " + "within 45 days of receipt."), + ("4. Delivery", "Delivery is made to the Buyer's nominated site. Risk passes on delivery."), + ("5. Termination", "Either party may end this Agreement by giving ninety (90) days prior written " + "notice. Notice is effective from the date of dispatch."), + ("6. Law and Jurisdiction", "This Agreement shall be governed by and construed in accordance with " + "the laws of England and Wales."), + ], + ) + return { + "contract_title": "Supply Agreement", + "party_a": "Sample Robotics Ltd.", + "party_b": "Demo Freight GmbH", + "effective_date": "2026-09-15", + "term_months": 12, + "notice_period_days": 90, + "payment_terms_days": 45, + "governing_law": "England and Wales", + } + + +# --------------------------------------------------------------------------- # +# Scanned-style invoice: rendered to an image, slightly rotated, image-only PDF +# --------------------------------------------------------------------------- # +def invoice_scanned(path: Path) -> dict: + import pypdfium2 as pdfium + from PIL import Image, ImageFilter + + tmp = io.BytesIO() + c = canvas.Canvas(tmp, pagesize=A4, invariant=1) + w, h = A4 + items = make_items(3, SERVICES, 50, 500) + sub, tax, total = totals(items, 0.09) + c.setFillColor(black) + c.setFont("Helvetica-Bold", 20) + c.drawString(60, h - 80, "Acme Example Services Pte. Ltd.") + c.setFont("Helvetica-Bold", 16) + c.drawString(60, h - 130, "INVOICE") + c.setFont("Helvetica", 12) + rows = [ + ("Invoice No.", "AES-5120"), + ("Date", "8 June 2026"), + ("Due", "8 July 2026"), + ("Billed to", "Northwind Parts Pte. Ltd."), + ] + y = h - 160 + for k, v in rows: + c.drawString(60, y, k) + c.drawString(180, y, v) + y -= 20 + y -= 20 + c.setFont("Helvetica-Bold", 12) + c.drawString(60, y, "Description") + c.drawRightString(w - 60, y, "Amount (SGD)") + c.setFont("Helvetica", 12) + for desc, qty, unit, amt in items: + y -= 22 + c.drawString(60, y, desc) + c.drawRightString(w - 60, y, money(amt)) + y -= 40 + for label, val in (("Subtotal", sub), ("GST 9%", tax), ("Total SGD", total)): + c.drawString(330, y, label) + c.drawRightString(w - 60, y, money(val)) + y -= 22 + c.save() + + page = pdfium.PdfDocument(tmp.getvalue())[0] + img = page.render(scale=150 / 72).to_pil().convert("L") + img = img.rotate(0.6, resample=Image.BICUBIC, expand=False, fillcolor=255) + import numpy as np + + noise_rng = np.random.default_rng(7) + noise = Image.fromarray(noise_rng.normal(128, 18, (img.height, img.width)).clip(0, 255).astype("uint8")) + img = Image.blend(img, noise, 0.08).filter(ImageFilter.GaussianBlur(0.4)) + img = img.point(lambda p: min(255, int(p * 1.04) + 6)) + + out = canvas.Canvas(str(path), pagesize=A4, invariant=1) + buf = io.BytesIO() + img.save(buf, format="PNG") + buf.seek(0) + out.drawImage(ImageReader(buf), 0, 0, width=w, height=h) + out.save() + return { + "invoice_number": "AES-5120", + "invoice_date": "2026-06-08", + "due_date": "2026-07-08", + "seller_name": "Acme Example Services Pte. Ltd.", + "buyer_name": "Northwind Parts Pte. Ltd.", + "currency": "SGD", + "subtotal": sub, + "tax": tax, + "total": total, + } + + +DOCS = [ + ("invoice-northwind-classic", "invoice", invoice_classic), + ("invoice-example-logistics-banner", "invoice", invoice_banner), + ("invoice-demo-freight-compact", "invoice", invoice_compact), + ("invoice-placeholder-two-page", "invoice", invoice_two_page), + ("contract-master-services", "contract", contract_services), + ("contract-supply", "contract", contract_supply), + ("invoice-acme-scanned", "invoice", invoice_scanned), +] + + +def main() -> None: + OUT.mkdir(exist_ok=True) + for name, schema, fn in DOCS: + fields = fn(OUT / f"{name}.pdf") + write_truth(name, schema, fields) + print(f"wrote samples/{name}.pdf") + + +if __name__ == "__main__": + main() diff --git a/src/document_extractor/__init__.py b/src/document_extractor/__init__.py index 48782cf..7882dde 100644 --- a/src/document_extractor/__init__.py +++ b/src/document_extractor/__init__.py @@ -1,6 +1,6 @@ """document-extractor: Turn PDF invoices and contracts into structured fields, each linked to its source. -Part of the Modelshifts showcase. Status: in development. +Part of the Modelshifts showcase. """ -__version__ = "0.0.1" +__version__ = "0.1.0" diff --git a/src/document_extractor/cli.py b/src/document_extractor/cli.py index 004e3eb..66f8b7a 100644 --- a/src/document_extractor/cli.py +++ b/src/document_extractor/cli.py @@ -1,14 +1,146 @@ -"""Command-line entry point (placeholder).""" +"""Command-line interface. + document-extractor extract FILE --schema schemas/invoice.yaml + document-extractor ask FILE "When does the notice period start?" + document-extractor batch DIR --schema invoice --out results.csv + document-extractor serve +""" + +from __future__ import annotations + +import argparse +import csv +import json import sys +import time +from pathlib import Path from . import __version__ +from .config import Settings +from .pdf import NoTextLayerError + + +def _settings() -> Settings: + settings = Settings.from_env() + settings.validate() + return settings + + +def cmd_extract(args: argparse.Namespace) -> int: + from .pipeline import run_extract + + settings = _settings() + result = run_extract(args.file, args.schema, settings) + text = json.dumps(result.to_json_dict(), indent=2, ensure_ascii=False) + if args.out: + Path(args.out).write_text(text + "\n") + print(f"wrote {args.out}", file=sys.stderr) + else: + print(text) + return 0 + + +def cmd_ask(args: argparse.Namespace) -> int: + from .pipeline import run_ask + + answer = run_ask(args.file, args.question, _settings()) + if args.json: + print(json.dumps(answer.model_dump(), indent=2, ensure_ascii=False)) + return 0 + print(answer.answer) + for c in answer.citations: + snippet = " ".join(c.text.split()) + print(f" [{c.n}] page {c.page}, box {[round(v, 1) for v in c.bbox]}: {snippet[:160]}") + return 0 + +def cmd_batch(args: argparse.Namespace) -> int: + from .pipeline import run_extract + from .schema import load_schema -def main() -> int: - print(f"document-extractor {__version__}: coming soon. See https://github.com/ModelShifts/document-extractor") + settings = _settings() + files = sorted(Path(args.dir).glob(args.glob)) + if not files: + print(f"no PDF files in {args.dir}", file=sys.stderr) + return 1 + field_names = load_schema(args.schema).field_names + results, rows = [], [] + for f in files: + t0 = time.time() + try: + r = run_extract(f, args.schema, settings) + except NoTextLayerError as exc: + print(f"skip {f.name}: {exc}", file=sys.stderr) + rows.append({"file": f.name, "error": str(exc)}) + continue + print(f"{f.name}: done in {time.time() - t0:.1f}s", file=sys.stderr) + results.append(r.to_json_dict()) + row = {"file": f.name, "error": ""} + for name in field_names: + fv = r.fields[name] + row[name] = "" if fv.value is None else fv.value + row[f"{name}_page"] = fv.source.page if fv.source else "" + rows.append(row) + out = Path(args.out) + if out.suffix.lower() == ".csv": + cols = ["file"] + [c for n in field_names for c in (n, f"{n}_page")] + ["error"] + with out.open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=cols, extrasaction="ignore") + writer.writeheader() + writer.writerows(rows) + else: + out.write_text(json.dumps(results, indent=2, ensure_ascii=False) + "\n") + print(f"wrote {out} ({len(results)} of {len(files)} files extracted)", file=sys.stderr) return 0 +def cmd_serve(args: argparse.Namespace) -> int: + from .webapp import launch + + launch(host=args.host, port=args.port, share=False) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(prog="document-extractor", description="Turn PDF invoices and contracts into " + "structured fields, each linked to its source.") + p.add_argument("--version", action="version", version=f"document-extractor {__version__}") + sub = p.add_subparsers(dest="command", required=True) + + e = sub.add_parser("extract", help="extract schema fields from one PDF as JSON") + e.add_argument("file") + e.add_argument("--schema", required=True, help="schema YAML path, or 'invoice' / 'contract'") + e.add_argument("--out", help="write JSON here instead of stdout") + e.set_defaults(func=cmd_extract) + + a = sub.add_parser("ask", help="ask a question about one PDF; the answer cites passages") + a.add_argument("file") + a.add_argument("question") + a.add_argument("--json", action="store_true", help="print the full answer as JSON") + a.set_defaults(func=cmd_ask) + + b = sub.add_parser("batch", help="extract every PDF in a folder to one CSV or JSON file") + b.add_argument("dir") + b.add_argument("--schema", required=True) + b.add_argument("--out", required=True, help="results.csv or results.json") + b.add_argument("--glob", default="*.pdf", help="which files to read, e.g. 'invoice-*.pdf' (default: *.pdf)") + b.set_defaults(func=cmd_batch) + + s = sub.add_parser("serve", help="start the local web demo (Gradio)") + s.add_argument("--host", default="127.0.0.1") + s.add_argument("--port", type=int, default=7860) + s.set_defaults(func=cmd_serve) + return p + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + try: + return args.func(args) + except (NoTextLayerError, ValueError, FileNotFoundError, RuntimeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + if __name__ == "__main__": sys.exit(main()) diff --git a/src/document_extractor/config.py b/src/document_extractor/config.py new file mode 100644 index 0000000..4458469 --- /dev/null +++ b/src/document_extractor/config.py @@ -0,0 +1,74 @@ +"""Settings read from environment variables (and a local ``.env`` file).""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path + +DEFAULT_HF_MODEL = "Qwen/Qwen2.5-7B-Instruct" +DEFAULT_HF_EMBED_MODEL = "BAAI/bge-m3" +DEFAULT_HF_OCR_MODEL = "florence-community/Florence-2-large" + + +def load_dotenv_if_present(path: str | os.PathLike | None = None) -> None: + """Load ``.env`` from the working directory without overriding real env vars.""" + try: + from dotenv import load_dotenv + except ImportError: # pragma: no cover - python-dotenv is a core dependency + return + load_dotenv(dotenv_path=path or Path.cwd() / ".env", override=False) + + +def _env(name: str, default: str = "") -> str: + value = os.environ.get(name, default) + return value.strip() if isinstance(value, str) else default + + +@dataclass +class Settings: + backend: str = "hf" + # api backend + api_base: str = "" + api_key: str = "" + model: str = "" + vision_model: str = "" + embed_model: str = "" + # hf backend + hf_model: str = DEFAULT_HF_MODEL + hf_embed_model: str = DEFAULT_HF_EMBED_MODEL + hf_ocr_model: str = DEFAULT_HF_OCR_MODEL + # shared + retriever: str = "auto" # auto | embed | bm25 + max_new_tokens: int = 512 + cache_dir: str = field(default_factory=lambda: _env("MS_CACHE_DIR", ".cache/results")) + + @classmethod + def from_env(cls) -> "Settings": + load_dotenv_if_present() + return cls( + backend=_env("MS_BACKEND", "hf").lower(), + api_base=_env("MS_API_BASE"), + api_key=_env("MS_API_KEY"), + model=_env("MS_MODEL"), + vision_model=_env("MS_VISION_MODEL"), + embed_model=_env("MS_EMBED_MODEL"), + hf_model=_env("MS_HF_MODEL", DEFAULT_HF_MODEL), + hf_embed_model=_env("MS_HF_EMBED_MODEL", DEFAULT_HF_EMBED_MODEL), + hf_ocr_model=_env("MS_HF_OCR_MODEL", DEFAULT_HF_OCR_MODEL), + retriever=_env("MS_RETRIEVER", "auto").lower(), + max_new_tokens=int(_env("MS_MAX_NEW_TOKENS", "512") or 512), + cache_dir=_env("MS_CACHE_DIR", ".cache/results"), + ) + + @property + def model_name(self) -> str: + return self.model if self.backend == "api" else self.hf_model + + def validate(self) -> None: + if self.backend not in {"api", "hf"}: + raise ValueError(f"MS_BACKEND must be 'api' or 'hf', got {self.backend!r}") + if self.backend == "api": + missing = [n for n, v in (("MS_API_BASE", self.api_base), ("MS_MODEL", self.model)) if not v] + if missing: + raise ValueError(f"MS_BACKEND=api needs {', '.join(missing)} (see .env.example)") diff --git a/src/document_extractor/extract.py b/src/document_extractor/extract.py new file mode 100644 index 0000000..0084560 --- /dev/null +++ b/src/document_extractor/extract.py @@ -0,0 +1,65 @@ +"""Schema-driven field extraction with a source for every value.""" + +from __future__ import annotations + +from pathlib import Path + +from .grounding import ground +from .llm import LLM, parse_json_object +from .models import Document, ExtractionResult, FieldValue +from .pdf import document_text +from .schema import Schema, coerce + +SYSTEM = ( + "You extract fields from business documents. Reply with a single JSON object and nothing else." +) + +MAX_DOC_CHARS = 12000 + + +def build_prompt(doc: Document, schema: Schema) -> list[dict[str, str]]: + lines = [f"- {f.name} ({f.type}): {f.description}" for f in schema.fields] + text = document_text(doc) + if len(text) > MAX_DOC_CHARS: + text = text[:MAX_DOC_CHARS] + "\n[...truncated]" + example = "{" + ", ".join(f'"{f.name}": "..."' for f in schema.fields[:2]) + ", ...}" + user = ( + f"Document type: {schema.name}. {schema.description}\n\n" + "Fields to extract:\n" + "\n".join(lines) + "\n\n" + "Rules:\n" + "- Copy each value exactly as it is written in the document: same spelling, number format and " + "date format. Do not convert or calculate.\n" + "- For amounts, copy only the number (for example 1,234.50), not the currency code.\n" + "- Use null if a field is not in the document.\n" + f"- Reply with one JSON object with exactly these keys: {', '.join(schema.field_names)}.\n" + f" Example shape: {example}\n\n" + "Document text (columns separated by |):\n<<<\n" + text + "\n>>>" + ) + return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}] + + +def extract(doc: Document, schema: Schema, llm: LLM, backend: str = "") -> ExtractionResult: + """Ask the model for the schema fields, convert types and attach sources.""" + reply = llm.complete(build_prompt(doc, schema)) + try: + data = parse_json_object(reply) + except ValueError: + data = {} + fields: dict[str, FieldValue] = {} + for spec in schema.fields: + raw = data.get(spec.name) + if isinstance(raw, (dict, list)): + raw = None + raw_text = None if raw is None else str(raw) + value = coerce(raw, spec.type) + source = None + if value is not None: + source = ground(doc, raw_text, value, spec.type, hint=f"{spec.name} {spec.description}") + fields[spec.name] = FieldValue(value=value, raw=raw_text, source=source) + return ExtractionResult( + file=Path(doc.path).name, + schema=schema.name, + backend=backend, + model=getattr(llm, "name", "unknown"), + fields=fields, + ) diff --git a/src/document_extractor/grounding.py b/src/document_extractor/grounding.py new file mode 100644 index 0000000..9d8b1a6 --- /dev/null +++ b/src/document_extractor/grounding.py @@ -0,0 +1,146 @@ +"""Link an extracted value back to the words on the page it came from. + +The model is asked to copy values exactly as written. We then look for the +best matching run of consecutive words on any page. If the model normalised a +value (for example ``2026-03-14`` instead of ``14 March 2026``), common +spellings of the typed value are tried as well. +""" + +from __future__ import annotations + +from difflib import SequenceMatcher +from typing import Any, Iterable + +from .models import Document, Source, Word +from .pdf import page_lines + +STRIP = " \t.,:;()[]{}\"'|*" +MIN_SCORE = 0.86 +MONTH_NAMES = ["January", "February", "March", "April", "May", "June", "July", "August", "September", + "October", "November", "December"] + + +def norm(token: str) -> str: + return token.strip(STRIP).lower() + + +def _key(tokens: Iterable[str]) -> str: + return " ".join(t for t in (norm(x) for x in tokens) if t) + + +def variants(raw: str | None, value: Any, ftype: str) -> list[str]: + """Candidate spellings to search for, most specific first.""" + out: list[str] = [] + if raw: + out.append(str(raw)) + if value is not None and ftype == "number": + x = float(value) + out += [f"{x:,.2f}", f"{x:.2f}", f"{x:,.2f}".replace(",", "_").replace(".", ",").replace("_", ".")] + elif value is not None and ftype == "date": + y, m, d = (int(v) for v in str(value).split("-")) + mon = MONTH_NAMES[m - 1] + out += [str(value), f"{d} {mon} {y}", f"{d:02d}/{m:02d}/{y}", f"{d:02d}.{m:02d}.{y}", f"{mon} {d}, {y}", + f"{d} {mon[:3]} {y}"] + elif value is not None and ftype == "integer": + out += [f"({value})", str(value)] + elif value is not None: + out.append(str(value)) + seen, uniq = set(), [] + for v in out: + if v and v not in seen: + seen.add(v) + uniq.append(v) + return uniq + + +def _union(words: list[Word]) -> tuple[float, float, float, float]: + return (round(min(w.bbox[0] for w in words), 2), round(min(w.bbox[1] for w in words), 2), + round(max(w.bbox[2] for w in words), 2), round(max(w.bbox[3] for w in words), 2)) + + +def _reading_order(doc: Document) -> list[Word]: + ordered: list[Word] = [] + for page in doc.pages: + for line in page_lines(page): + ordered.extend(line) + return ordered + + +STOP = {"the", "a", "an", "of", "and", "or", "to", "in", "on", "for", "by", "is", "as", "at", "this", "that", + "which", "who", "with", "from", "be", "it", "its", "date", "name"} + + +def _stems(text: str) -> set[str]: + out = set() + for t in text.replace("_", " ").split(): + t = norm(t) + if len(t) > 3 and t.endswith("s"): + t = t[:-1] + if len(t) > 2 and t not in STOP: + out.add(t[:5]) + return out + + +def _context_score(span: list[Word], words: list[Word], hint: set[str]) -> int: + """How many hint words appear just before the span on the same page (label proximity).""" + if not hint: + return 0 + first = words.index(span[0]) + before = [w for w in words[max(0, first - 12) : first] if w.page == span[0].page] + return len(hint & _stems(" ".join(w.text for w in before))) + + +def find_span( + doc: Document, text: str, words: list[Word] | None = None, hint: str = "" +) -> tuple[float, list[Word]] | None: + """Best matching run of words for ``text``: (score, words) or ``None``. + + When the same text appears more than once, the copy whose preceding words + share the most terms with ``hint`` (field name and description) wins. + """ + target = _key(text.split()) + if not target: + return None + words = words if words is not None else _reading_order(doc) + n = len(target.split()) + hint_stems = _stems(hint) + exact: list[list[Word]] = [] + best: tuple[float, list[Word]] | None = None + for size in sorted({n, max(1, n - 1), n + 1}): + for i in range(0, len(words) - size + 1): + span = words[i : i + size] + if len({w.page for w in span}) > 1: + continue + cand = _key(w.text for w in span) + if not cand: + continue + if cand == target: + score = 1.0 + elif size == 1 and n == 1 and target in cand and len(target) >= 4: + score = 0.95 # value printed inside a longer token, e.g. "SGD5,496.68" + else: + score = SequenceMatcher(None, cand, target).ratio() + if score == 1.0: + exact.append(span) + elif best is None or score > best[0] + 1e-9: + best = (score, span) + if exact: + ranked = sorted(enumerate(exact), key=lambda e: (-_context_score(e[1], words, hint_stems), e[0])) + return 1.0, ranked[0][1] + return best + + +def ground(doc: Document, raw: str | None, value: Any, ftype: str = "string", hint: str = "") -> Source | None: + """Return the page, box and text of the best match, or ``None`` if nothing is close enough.""" + words = _reading_order(doc) + best: tuple[float, list[Word]] | None = None + for candidate in variants(raw, value, ftype): + found = find_span(doc, candidate, words, hint) + if found and (best is None or found[0] > best[0]): + best = found + if best and best[0] == 1.0: + break + if not best or best[0] < MIN_SCORE: + return None + span = best[1] + return Source(page=span[0].page, bbox=_union(span), text=" ".join(w.text for w in span)) diff --git a/src/document_extractor/llm.py b/src/document_extractor/llm.py new file mode 100644 index 0000000..3f5ce76 --- /dev/null +++ b/src/document_extractor/llm.py @@ -0,0 +1,133 @@ +"""Chat model backends. + +``api``: any OpenAI-compatible ``/v1/chat/completions`` endpoint. +``hf``: a local ``transformers`` causal LM (default Qwen/Qwen2.5-7B-Instruct). + +Both expose ``complete(messages) -> str``. Decoding is greedy (temperature 0). +""" + +from __future__ import annotations + +import json +import re +from typing import Any, Protocol + +from .config import Settings + +Message = dict[str, str] + + +class LLM(Protocol): + name: str + + def complete(self, messages: list[Message], max_new_tokens: int | None = None) -> str: ... + + +class APIChat: + def __init__(self, settings: Settings): + try: + from openai import OpenAI + except ImportError as exc: # pragma: no cover + raise RuntimeError("MS_BACKEND=api needs the 'api' extra: pip install -e '.[api]'") from exc + self.client = OpenAI(base_url=settings.api_base, api_key=settings.api_key or "not-needed") + self.name = settings.model + self.max_new_tokens = settings.max_new_tokens + + def complete(self, messages: list[Message], max_new_tokens: int | None = None) -> str: + response = self.client.chat.completions.create( + model=self.name, + messages=messages, + temperature=0, + max_tokens=max_new_tokens or self.max_new_tokens, + ) + return response.choices[0].message.content or "" + + +class HFChat: + def __init__(self, settings: Settings): + try: + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + except ImportError as exc: # pragma: no cover + raise RuntimeError("MS_BACKEND=hf needs the 'hf' extra: pip install -e '.[hf]'") from exc + self._torch = torch + self.name = settings.hf_model + self.max_new_tokens = settings.max_new_tokens + self.device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 + self.tokenizer = AutoTokenizer.from_pretrained(self.name) + self.model = AutoModelForCausalLM.from_pretrained(self.name, dtype=dtype).to(self.device).eval() + + def complete(self, messages: list[Message], max_new_tokens: int | None = None) -> str: + prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device) + with self._torch.no_grad(): + out = self.model.generate( + **inputs, + max_new_tokens=max_new_tokens or self.max_new_tokens, + do_sample=False, + temperature=None, + top_p=None, + top_k=None, + pad_token_id=self.tokenizer.eos_token_id, + ) + return self.tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True) + + +_CACHE: dict[tuple, Any] = {} + + +def get_llm(settings: Settings) -> LLM: + """Build (and memoise) the chat model for the configured backend.""" + settings.validate() + key = ("llm", settings.backend, settings.model_name, settings.api_base) + if key not in _CACHE: + _CACHE[key] = APIChat(settings) if settings.backend == "api" else HFChat(settings) + return _CACHE[key] + + +def get_ocr(settings: Settings): + """Return an OCR engine for the backend, or ``None`` if none is configured.""" + key = ("ocr", settings.backend, settings.hf_ocr_model, settings.vision_model) + if key in _CACHE: + return _CACHE[key] + if settings.backend == "hf": + from .ocr import Florence2OCR + + engine = Florence2OCR(settings.hf_ocr_model) + elif settings.vision_model: + from .ocr import APIVisionOCR + + chat = get_llm(settings) + engine = APIVisionOCR(chat.client, settings.vision_model) # type: ignore[attr-defined] + else: + engine = None + _CACHE[key] = engine + return engine + + +def parse_json_object(text: str) -> dict[str, Any]: + """Pull the first JSON object out of a model reply (tolerates code fences and chatter).""" + text = re.sub(r"```(?:json)?", "", text) + start = text.find("{") + if start < 0: + raise ValueError(f"model reply has no JSON object: {text[:200]!r}") + depth, in_str, esc = 0, False, False + for i in range(start, len(text)): + ch = text[i] + if in_str: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_str = False + elif ch == '"': + in_str = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return json.loads(text[start : i + 1]) + raise ValueError(f"model reply has an unterminated JSON object: {text[:200]!r}") diff --git a/src/document_extractor/models.py b/src/document_extractor/models.py new file mode 100644 index 0000000..1af62e0 --- /dev/null +++ b/src/document_extractor/models.py @@ -0,0 +1,76 @@ +"""Data models shared by the extractor, the Q&A module and the web demo.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + +BBox = tuple[float, float, float, float] # x0, top, x1, bottom in PDF points, origin top-left + + +class Word(BaseModel): + text: str + page: int # 1-based + bbox: BBox + + +class Page(BaseModel): + number: int # 1-based + width: float + height: float + words: list[Word] = Field(default_factory=list) + ocr: bool = False # True when the words came from OCR instead of the PDF text layer + + +class Document(BaseModel): + path: str + pages: list[Page] + + @property + def words(self) -> list[Word]: + return [w for p in self.pages for w in p.words] + + +class Source(BaseModel): + page: int + bbox: BBox + text: str + + +class FieldValue(BaseModel): + value: Any = None + raw: str | None = None # the text the model returned before type conversion + source: Source | None = None + + +class ExtractionResult(BaseModel): + file: str + schema_name: str = Field(alias="schema") + backend: str + model: str + fields: dict[str, FieldValue] + + model_config = {"populate_by_name": True} + + def values(self) -> dict[str, Any]: + return {k: v.value for k, v in self.fields.items()} + + def to_json_dict(self) -> dict[str, Any]: + return self.model_dump(by_alias=True) + + +class Citation(BaseModel): + n: int + page: int + bbox: BBox + text: str + + +class Answer(BaseModel): + question: str + answer: str + citations: list[Citation] + backend: str + model: str + retriever: str diff --git a/src/document_extractor/ocr.py b/src/document_extractor/ocr.py new file mode 100644 index 0000000..290e28d --- /dev/null +++ b/src/document_extractor/ocr.py @@ -0,0 +1,132 @@ +"""OCR for pages without a text layer. + +* ``hf`` backend: Florence-2 (```` task) gives text lines with boxes. +* ``api`` backend: a vision model behind the same OpenAI-compatible endpoint, if + ``MS_VISION_MODEL`` is set. It is asked for text lines with boxes on a 0-1000 + grid. Box quality depends on the model. + +Both return :class:`~document_extractor.models.Word` objects in PDF points. +""" + +from __future__ import annotations + +import base64 +import io +import json +import re +from typing import Protocol + +from PIL import Image + +from .models import Word + + +class OCREngine(Protocol): + name: str + + def recognise(self, image: Image.Image, page: int, scale: float) -> list[Word]: ... + + +def split_line(text: str, box: tuple[float, float, float, float], page: int) -> list[Word]: + """Split one OCR line into words, sharing the line box out by character count.""" + x0, y0, x1, y1 = box + tokens = text.split() + if not tokens: + return [] + total = sum(len(t) for t in tokens) + (len(tokens) - 1) + width = max(x1 - x0, 1e-6) + words, cursor = [], 0 + for t in tokens: + a = x0 + width * cursor / total + b = x0 + width * (cursor + len(t)) / total + words.append(Word(text=t, page=page, bbox=(a, y0, b, y1))) + cursor += len(t) + 1 + return words + + +class Florence2OCR: + """Florence-2 OCR via ``transformers`` (runs on CPU or GPU).""" + + def __init__(self, model_id: str): + import torch + from transformers import AutoProcessor, Florence2ForConditionalGeneration + + self.name = model_id + self._torch = torch + self.device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float16 if self.device == "cuda" else torch.float32 + self.model = Florence2ForConditionalGeneration.from_pretrained(model_id, dtype=dtype).to(self.device).eval() + self.processor = AutoProcessor.from_pretrained(model_id) + self.dtype = dtype + + def recognise(self, image: Image.Image, page: int, scale: float) -> list[Word]: + task = "" + inputs = self.processor(text=task, images=image.convert("RGB"), return_tensors="pt").to(self.device, self.dtype) + with self._torch.no_grad(): + generated = self.model.generate(**inputs, max_new_tokens=1024, num_beams=3, do_sample=False) + text = self.processor.batch_decode(generated, skip_special_tokens=False)[0] + parsed = self.processor.post_process_generation(text, task=task, image_size=image.size)[task] + words: list[Word] = [] + for label, quad in zip(parsed.get("labels", []), parsed.get("quad_boxes", [])): + xs, ys = quad[0::2], quad[1::2] + box = (min(xs) / scale, min(ys) / scale, max(xs) / scale, max(ys) / scale) + words.extend(split_line(label.replace("", "").strip(), box, page)) + return words + + +VISION_PROMPT = ( + "Read all text on this page. Reply with JSON only: a list of objects " + '{"text": "", "box": [x0, y0, x1, y1]} where the box is the line position ' + "on a 0-1000 grid (0,0 = top-left, 1000,1000 = bottom-right). Keep the reading order." +) + + +class APIVisionOCR: + """OCR with a vision-capable chat model on an OpenAI-compatible endpoint.""" + + def __init__(self, client, model: str): + self.client = client + self.name = model + self.model = model + + def recognise(self, image: Image.Image, page: int, scale: float) -> list[Word]: + buf = io.BytesIO() + image.convert("RGB").save(buf, format="PNG") + url = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() + response = self.client.chat.completions.create( + model=self.model, + temperature=0, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": VISION_PROMPT}, + {"type": "image_url", "image_url": {"url": url}}, + ], + } + ], + ) + return parse_vision_lines(response.choices[0].message.content or "", image.size, page, scale) + + +def parse_vision_lines(content: str, size: tuple[int, int], page: int, scale: float) -> list[Word]: + match = re.search(r"\[.*\]", content, re.S) + if not match: + return [] + try: + items = json.loads(match.group(0)) + except json.JSONDecodeError: + return [] + w_pts, h_pts = size[0] / scale, size[1] / scale + words: list[Word] = [] + for item in items: + if not isinstance(item, dict) or not item.get("text"): + continue + box = item.get("box") or [0, 0, 1000, 1000] + try: + x0, y0, x1, y1 = (float(v) for v in box[:4]) + except (TypeError, ValueError): + continue + pts = (x0 / 1000 * w_pts, y0 / 1000 * h_pts, x1 / 1000 * w_pts, y1 / 1000 * h_pts) + words.extend(split_line(str(item["text"]), pts, page)) + return words diff --git a/src/document_extractor/pdf.py b/src/document_extractor/pdf.py new file mode 100644 index 0000000..a5f3410 --- /dev/null +++ b/src/document_extractor/pdf.py @@ -0,0 +1,110 @@ +"""PDF loading: words with page numbers and bounding boxes, plus page images. + +Text-layer PDFs are read with pdfplumber. Pages without a text layer (scans) +are rendered with pypdfium2 and passed to an OCR engine (see ``ocr.py``). +Coordinates are PDF points with the origin at the top-left of the page. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import pdfplumber +import pypdfium2 as pdfium +from PIL import Image + +from .models import Document, Page, Word + +if TYPE_CHECKING: # pragma: no cover + from .ocr import OCREngine + +MIN_TEXT_WORDS = 3 # fewer words than this on a page with an image -> treat as a scan +OCR_SCALE = 2.0 # render scale for OCR (2.0 = 144 dpi) + + +def render_page(path: str | Path, page_number: int, scale: float = 2.0) -> Image.Image: + """Render one page (1-based) to an RGB image. ``scale`` is pixels per PDF point.""" + pdf = pdfium.PdfDocument(str(path)) + try: + return pdf[page_number - 1].render(scale=scale).to_pil().convert("RGB") + finally: + pdf.close() + + +def load_pdf(path: str | Path, ocr: "OCREngine | None" = None) -> Document: + """Read every page of ``path`` into a :class:`Document`. + + If a page has no usable text layer, ``ocr`` is used to recognise it. When no + OCR engine is given, :class:`NoTextLayerError` explains how to enable one. + """ + path = Path(path) + pages: list[Page] = [] + with pdfplumber.open(str(path)) as pdf: + for i, p in enumerate(pdf.pages, start=1): + words = [ + Word(text=w["text"], page=i, bbox=(float(w["x0"]), float(w["top"]), float(w["x1"]), float(w["bottom"]))) + for w in p.extract_words(keep_blank_chars=False, use_text_flow=False) + ] + page = Page(number=i, width=float(p.width), height=float(p.height), words=words) + if len(words) < MIN_TEXT_WORDS and (p.images or not words): + if ocr is None: + raise NoTextLayerError(path.name, i) + image = render_page(path, i, scale=OCR_SCALE) + page.words = ocr.recognise(image, page=i, scale=OCR_SCALE) + page.ocr = True + pages.append(page) + return Document(path=str(path), pages=pages) + + +def has_text_layer(path: str | Path) -> bool: + with pdfplumber.open(str(path)) as pdf: + return all(len(p.extract_words()) >= MIN_TEXT_WORDS for p in pdf.pages) + + +class NoTextLayerError(RuntimeError): + def __init__(self, name: str, page: int): + super().__init__( + f"{name}, page {page}: this page has no text layer (it looks like a scan) and no OCR engine is " + "configured. Use MS_BACKEND=hf (Florence-2 OCR) or set MS_VISION_MODEL for MS_BACKEND=api." + ) + self.page = page + + +def page_lines(page: Page, y_tolerance: float = 3.0) -> list[list[Word]]: + """Group a page's words into visual lines, top to bottom, left to right.""" + lines: list[list[Word]] = [] + for w in sorted(page.words, key=lambda w: (w.bbox[1], w.bbox[0])): + mid = (w.bbox[1] + w.bbox[3]) / 2 + for line in lines: + ref = line[0] + if abs(mid - (ref.bbox[1] + ref.bbox[3]) / 2) <= y_tolerance: + line.append(w) + break + else: + lines.append([w]) + for line in lines: + line.sort(key=lambda w: w.bbox[0]) + lines.sort(key=lambda ln: min(w.bbox[1] for w in ln)) + return lines + + +def line_text(line: list[Word], gap: float = 14.0) -> str: + """Join a line's words; wide horizontal gaps become ``" | "`` so table columns stay apart.""" + out = "" + prev: Word | None = None + for w in line: + if prev is not None: + out += " | " if w.bbox[0] - prev.bbox[2] > gap else " " + out += w.text + prev = w + return out + + +def document_text(doc: Document) -> str: + """Plain-text view of the document used in prompts, with page markers.""" + parts = [] + for page in doc.pages: + parts.append(f"[Page {page.number}]") + parts.extend(line_text(line) for line in page_lines(page)) + return "\n".join(parts) diff --git a/src/document_extractor/pipeline.py b/src/document_extractor/pipeline.py new file mode 100644 index 0000000..20d0cdb --- /dev/null +++ b/src/document_extractor/pipeline.py @@ -0,0 +1,68 @@ +"""High-level helpers used by the CLI and the web demo.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +from .config import Settings +from .extract import extract +from .llm import get_llm, get_ocr +from .models import Answer, Document, ExtractionResult +from .pdf import NoTextLayerError, load_pdf +from .qa import ask +from .retrieval import get_retriever +from .schema import load_schema + +_DOCS: dict[tuple[str, str], Document] = {} + + +def file_digest(path: str | Path) -> str: + return hashlib.sha256(Path(path).read_bytes()).hexdigest()[:16] + + +def load_document(path: str | Path, settings: Settings) -> Document: + """Load a PDF, running OCR only when a page has no text layer. Results are memoised per file.""" + key = (file_digest(path), settings.backend) + if key in _DOCS: + return _DOCS[key] + try: + doc = load_pdf(path) + except NoTextLayerError: + ocr = get_ocr(settings) + if ocr is None: + raise + doc = load_pdf(path, ocr=ocr) + _DOCS[key] = doc + return doc + + +def run_extract(path: str | Path, schema: str | Path, settings: Settings) -> ExtractionResult: + doc = load_document(path, settings) + return extract(doc, load_schema(schema), get_llm(settings), backend=settings.backend) + + +def run_ask(path: str | Path, question: str, settings: Settings) -> Answer: + doc = load_document(path, settings) + return ask(doc, question, get_llm(settings), get_retriever(settings), backend=settings.backend) + + +class ResultCache: + """Small JSON cache keyed by file content, task and model. Used by the web demo.""" + + def __init__(self, root: str | Path): + self.root = Path(root) + + def _path(self, path: str | Path, kind: str, key: str, settings: Settings) -> Path: + tag = hashlib.sha256(f"{kind}|{key}|{settings.backend}|{settings.model_name}".encode()).hexdigest()[:16] + return self.root / f"{file_digest(path)}-{kind}-{tag}.json" + + def get(self, path, kind: str, key: str, settings: Settings) -> dict | None: + p = self._path(path, kind, key, settings) + return json.loads(p.read_text()) if p.exists() else None + + def put(self, path, kind: str, key: str, settings: Settings, data: dict) -> None: + p = self._path(path, kind, key, settings) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(data, indent=2)) diff --git a/src/document_extractor/qa.py b/src/document_extractor/qa.py new file mode 100644 index 0000000..ff2a86c --- /dev/null +++ b/src/document_extractor/qa.py @@ -0,0 +1,79 @@ +"""Question answering over one document, with cited passages.""" + +from __future__ import annotations + +import re + +from .llm import LLM +from .models import Answer, Citation, Document +from .retrieval import Passage, Retriever, chunk + +SYSTEM = ( + "You answer questions about a business document using only the numbered passages you are given. " + "Write the answer as one short sentence, then put the number of the passage you used in square " + "brackets at the end. Example: The term is 24 months [2]. " + "Only if none of the passages answers the question, reply: Not found in the document." +) + + +def build_prompt(question: str, passages: list[Passage]) -> list[dict[str, str]]: + blocks = [f"[{i}] (page {p.page})\n{p.text}" for i, p in enumerate(passages, start=1)] + user = ( + "Passages:\n\n" + "\n\n".join(blocks) + f"\n\nQuestion: {question}\n" + "Answer in one sentence, ending with the passage number in brackets." + ) + return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}] + + +def _content_tokens(text: str) -> set[str]: + text = re.sub(r"\[\d+\]", " ", text.lower()) + return {t for t in re.findall(r"[a-z0-9][a-z0-9,.\-]*[a-z0-9]|\d", text) if any(c.isdigit() for c in t) or len(t) > 3} + + +def _repair_citations(reply: str, passages: list[Passage], cited: list[int]) -> tuple[list[int], dict[int, int]]: + """Small models sometimes cite the wrong number. If another retrieved passage shares at + least two more content words with the answer than the cited one, cite that one instead.""" + words = _content_tokens(reply) + if not words: + return cited, {} + overlap = [len(words & _content_tokens(p.text)) for p in passages] + fixed, moved = [], {} + for i in cited: + best = max(range(len(passages)), key=lambda j: overlap[j]) + if overlap[best] >= overlap[i - 1] + 2: + moved[i] = best + 1 + i = best + 1 + if i not in fixed: + fixed.append(i) + return fixed, moved + + +def ask(doc: Document, question: str, llm: LLM, retriever: Retriever, k: int = 4, backend: str = "") -> Answer: + passages = retriever.search(question, chunk(doc), k) + reply = llm.complete(build_prompt(question, passages), max_new_tokens=200).strip() + if reply.lower().startswith("not found"): + reply = "Not found in the document." + cited = [] + for n in re.findall(r"\[(\d+)\]", reply): + i = int(n) + if 1 <= i <= len(passages) and i not in cited: + cited.append(i) + if not cited and passages and not reply.lower().startswith("not found"): + cited = [1] # model forgot to cite: point at the top-ranked passage + cited, moved = _repair_citations(reply, passages, cited) + if moved: # keep the numbers in the answer text in step with the citations + reply = re.sub(r"\[(\d+)\]", lambda m: f"[{moved.get(int(m.group(1)), int(m.group(1)))}]", reply) + if reply == "Not found in the document.": + cited = [] + citations = [ + Citation(n=i, page=passages[i - 1].page, bbox=passages[i - 1].bbox, text=passages[i - 1].text) + for i in cited + ] + return Answer( + question=question, + answer=reply, + citations=citations, + backend=backend, + model=getattr(llm, "name", "unknown"), + retriever=retriever.name, + ) diff --git a/src/document_extractor/retrieval.py b/src/document_extractor/retrieval.py new file mode 100644 index 0000000..347eec4 --- /dev/null +++ b/src/document_extractor/retrieval.py @@ -0,0 +1,166 @@ +"""Split a document into passages and retrieve the ones relevant to a question. + +Retrievers: + +* ``embed`` with the ``hf`` backend: BAAI/bge-m3 dense embeddings (sentence-transformers). +* ``embed`` with the ``api`` backend: the endpoint's ``/v1/embeddings`` with ``MS_EMBED_MODEL``. +* ``bm25``: keyword ranking (rank-bm25). No model download, used by the tests and as a + fallback when no embedding model is configured. + +``MS_RETRIEVER=auto`` (default) picks ``embed`` when an embedding model is +available for the backend and ``bm25`` otherwise. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from statistics import median +from typing import Protocol + +from .config import Settings +from .models import Document, Word +from .pdf import line_text, page_lines + + +@dataclass +class Passage: + page: int + text: str + words: list[Word] + + @property + def bbox(self) -> tuple[float, float, float, float]: + ws = self.words + return (round(min(w.bbox[0] for w in ws), 2), round(min(w.bbox[1] for w in ws), 2), + round(max(w.bbox[2] for w in ws), 2), round(max(w.bbox[3] for w in ws), 2)) + + +HEADING = re.compile(r"^(\d+(\.\d+)*\.?|[A-Z][A-Z ]{3,})\s+[A-Z]") + + +def chunk(doc: Document, max_chars: int = 600) -> list[Passage]: + """Paragraph-like passages. + + A new passage starts at a vertical gap larger than a line height, at a numbered + or upper-case heading (``7. Governing Law``), or when ``max_chars`` is reached. + """ + passages: list[Passage] = [] + for page in doc.pages: + lines = page_lines(page) + if not lines: + continue + heights = [max(w.bbox[3] for w in ln) - min(w.bbox[1] for w in ln) for ln in lines] + h = median(heights) or 10.0 + cur_words: list[Word] = [] + cur_text: list[str] = [] + prev_bottom = None + for ln in lines: + top = min(w.bbox[1] for w in ln) + text = line_text(ln) + gap = None if prev_bottom is None else top - prev_bottom + too_long = sum(len(t) for t in cur_text) + len(text) > max_chars + heading = bool(HEADING.match(text)) + if cur_words and ((gap is not None and gap > 0.9 * h) or too_long or heading): + passages.append(Passage(page.number, "\n".join(cur_text), cur_words)) + cur_words, cur_text = [], [] + cur_words = cur_words + ln + cur_text.append(text) + prev_bottom = max(w.bbox[3] for w in ln) + if cur_words: + passages.append(Passage(page.number, "\n".join(cur_text), cur_words)) + return passages + + +class Retriever(Protocol): + name: str + + def search(self, query: str, passages: list[Passage], k: int) -> list[Passage]: ... + + +def _tokens(text: str) -> list[str]: + return re.findall(r"[a-z0-9]+", text.lower()) + + +class BM25Retriever: + name = "bm25" + + def search(self, query: str, passages: list[Passage], k: int) -> list[Passage]: + from rank_bm25 import BM25Okapi + + if not passages: + return [] + bm25 = BM25Okapi([_tokens(p.text) or ["_"] for p in passages]) + scores = bm25.get_scores(_tokens(query)) + order = sorted(range(len(passages)), key=lambda i: -scores[i]) + return [passages[i] for i in order[:k]] + + +class _DenseRetriever: + name = "embed" + + def embed(self, texts: list[str]) -> list[list[float]]: # pragma: no cover - abstract + raise NotImplementedError + + def search(self, query: str, passages: list[Passage], k: int) -> list[Passage]: + import numpy as np + + if not passages: + return [] + vecs = np.asarray(self.embed([query] + [p.text for p in passages]), dtype="float32") + vecs /= np.linalg.norm(vecs, axis=1, keepdims=True) + 1e-12 + scores = vecs[1:] @ vecs[0] + order = np.argsort(-scores)[:k] + return [passages[int(i)] for i in order] + + +class BGEM3Retriever(_DenseRetriever): + def __init__(self, model_id: str): + from sentence_transformers import SentenceTransformer + + self.name = f"embed:{model_id}" + self.model = SentenceTransformer(model_id, device="cpu") + + def embed(self, texts: list[str]) -> list[list[float]]: + return self.model.encode(texts, normalize_embeddings=True, batch_size=16).tolist() + + +class APIEmbeddingRetriever(_DenseRetriever): + def __init__(self, client, model: str): + self.name = f"embed:{model}" + self.client = client + self.model = model + + def embed(self, texts: list[str]) -> list[list[float]]: + import numpy as np + + response = self.client.embeddings.create(model=self.model, input=texts) + out = [] + for d in response.data: + vec = np.asarray(d.embedding, dtype="float32") + if vec.ndim == 2: # some local servers return one vector per token: mean-pool them + vec = vec.mean(axis=0) + out.append(vec.tolist()) + return out + + +_RETRIEVERS: dict[tuple, Retriever] = {} + + +def get_retriever(settings: Settings) -> Retriever: + mode = settings.retriever + if mode not in {"auto", "embed", "bm25"}: + raise ValueError("MS_RETRIEVER must be auto, embed or bm25") + if mode == "bm25" or (mode == "auto" and settings.backend == "api" and not settings.embed_model): + return BM25Retriever() + if settings.backend == "api" and not settings.embed_model: + raise ValueError("MS_RETRIEVER=embed with MS_BACKEND=api needs MS_EMBED_MODEL") + key = (settings.backend, settings.hf_embed_model, settings.embed_model, settings.api_base) + if key not in _RETRIEVERS: + if settings.backend == "hf": + _RETRIEVERS[key] = BGEM3Retriever(settings.hf_embed_model) + else: + from .llm import get_llm + + _RETRIEVERS[key] = APIEmbeddingRetriever(get_llm(settings).client, settings.embed_model) # type: ignore[attr-defined] + return _RETRIEVERS[key] diff --git a/src/document_extractor/schema.py b/src/document_extractor/schema.py new file mode 100644 index 0000000..0dc07ec --- /dev/null +++ b/src/document_extractor/schema.py @@ -0,0 +1,160 @@ +"""YAML field schemas and type conversion. + +A schema file looks like:: + + name: invoice + description: A commercial invoice. + fields: + invoice_number: {type: string, description: The invoice identifier} + total: {type: number, description: The final amount due} + +Supported types: ``string``, ``date`` (ISO ``YYYY-MM-DD`` output), ``number`` +(float) and ``integer``. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from datetime import date +from pathlib import Path +from typing import Any + +import yaml + +TYPES = {"string", "date", "number", "integer"} +PACKAGE_SCHEMAS = Path(__file__).resolve().parent / "schemas" # wheel installs +REPO_SCHEMAS = Path(__file__).resolve().parents[2] / "schemas" # source checkout / editable install + + +@dataclass +class FieldSpec: + name: str + type: str + description: str + + +@dataclass +class Schema: + name: str + description: str + fields: list[FieldSpec] + + @property + def field_names(self) -> list[str]: + return [f.name for f in self.fields] + + +def resolve_schema_path(name_or_path: str | Path) -> Path: + """Accept a path, or a bare name like ``invoice`` for a bundled schema.""" + p = Path(name_or_path) + if p.exists(): + return p + stem = Path(name_or_path).stem + for folder in (Path("schemas"), REPO_SCHEMAS, PACKAGE_SCHEMAS): + candidate = folder / f"{stem}.yaml" + if candidate.exists(): + return candidate + raise FileNotFoundError(f"schema not found: {name_or_path}") + + +def load_schema(name_or_path: str | Path) -> Schema: + path = resolve_schema_path(name_or_path) + data = yaml.safe_load(path.read_text()) + fields = [] + for name, spec in (data.get("fields") or {}).items(): + spec = spec or {} + ftype = str(spec.get("type", "string")) + if ftype not in TYPES: + raise ValueError(f"{path}: field {name!r} has unknown type {ftype!r} (use one of {sorted(TYPES)})") + fields.append(FieldSpec(name=name, type=ftype, description=str(spec.get("description", "")))) + if not fields: + raise ValueError(f"{path}: schema has no fields") + return Schema(name=str(data.get("name", path.stem)), description=str(data.get("description", "")), fields=fields) + + +# --------------------------------------------------------------------------- # +# Type conversion +# --------------------------------------------------------------------------- # +MONTHS = {m: i for i, m in enumerate( + ["january", "february", "march", "april", "may", "june", "july", "august", "september", "october", + "november", "december"], start=1)} +MONTHS.update({k[:3]: v for k, v in list(MONTHS.items())}) +WORD_NUMBERS = { + "one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six": 6, "seven": 7, "eight": 8, "nine": 9, + "ten": 10, "eleven": 11, "twelve": 12, "fourteen": 14, "fifteen": 15, "twenty": 20, "thirty": 30, + "forty": 40, "forty-five": 45, "sixty": 60, "ninety": 90, +} + + +def parse_date(text: str) -> str | None: + s = text.strip().lower().replace(",", " ") + if m := re.search(r"\b(\d{4})-(\d{1,2})-(\d{1,2})\b", s): + y, mo, d = map(int, m.groups()) + elif m := re.search(r"\b(\d{1,2})[./](\d{1,2})[./](\d{4})\b", s): + d, mo, y = map(int, m.groups()) # day-first, as on the sample documents + if mo > 12 and d <= 12: + d, mo = mo, d + elif m := re.search(r"\b(\d{1,2})(?:st|nd|rd|th)?\s+([a-z]+)\s+(\d{4})\b", s): + d, mon, y = int(m.group(1)), m.group(2), int(m.group(3)) + if mon not in MONTHS: + return None + mo = MONTHS[mon] + elif m := re.search(r"\b([a-z]+)\s+(\d{1,2})(?:st|nd|rd|th)?\s+(\d{4})\b", s): + mon, d, y = m.group(1), int(m.group(2)), int(m.group(3)) + if mon not in MONTHS: + return None + mo = MONTHS[mon] + else: + return None + try: + return date(y, mo, d).isoformat() + except ValueError: + return None + + +def parse_number(text: str) -> float | None: + s = re.sub(r"[^\d,.\-]", "", str(text)) + if not re.search(r"\d", s): + return None + if "," in s and "." in s: + if s.rfind(",") > s.rfind("."): # 1.234,56 + s = s.replace(".", "").replace(",", ".") + else: # 1,234.56 + s = s.replace(",", "") + elif "," in s: + head, _, tail = s.rpartition(",") + s = s.replace(",", "") if len(tail) == 3 else f"{head.replace(',', '')}.{tail}" + try: + return float(s) + except ValueError: + return None + + +def parse_integer(text: str) -> int | None: + s = str(text).lower() + if m := re.search(r"\(\s*(\d+)\s*\)", s): + return int(m.group(1)) + if m := re.search(r"-?\d[\d,]*", s): + return int(m.group(0).replace(",", "")) + for word, n in sorted(WORD_NUMBERS.items(), key=lambda kv: -len(kv[0])): + if re.search(rf"\b{word}\b", s): + return n + return None + + +def coerce(value: Any, ftype: str) -> Any: + """Convert a raw model value to the schema type. Returns ``None`` when it cannot.""" + if value is None: + return None + if isinstance(value, str) and value.strip().lower() in {"", "null", "none", "n/a", "not found"}: + return None + if ftype == "string": + return str(value).strip() + if ftype == "date": + return parse_date(str(value)) + if ftype == "number": + return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else parse_number(value) + if ftype == "integer": + return int(value) if isinstance(value, int) and not isinstance(value, bool) else parse_integer(str(value)) + return value diff --git a/src/document_extractor/webapp.py b/src/document_extractor/webapp.py new file mode 100644 index 0000000..d93b7fa --- /dev/null +++ b/src/document_extractor/webapp.py @@ -0,0 +1,286 @@ +"""Local web demo (Gradio): pick or upload a PDF, extract fields, see sources, ask questions.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFont + +from . import __version__ +from .config import Settings +from .models import Answer, ExtractionResult +from .pdf import NoTextLayerError, render_page +from .pipeline import ResultCache, run_ask, run_extract + +import pypdfium2 as pdfium + +RENDER_SCALE = 1.6 +FIELD_COLOUR = (138, 58, 36) # rust +CITE_COLOUR = (28, 26, 23) # ink +CITE_FILL = (231, 196, 92, 90) # soft amber +PAGE_GAP = 18 + + +def samples_dir() -> Path: + for d in (Path.cwd() / "samples", Path(__file__).resolve().parents[2] / "samples"): + if d.is_dir(): + return d + return Path.cwd() / "samples" + + +def list_samples() -> list[str]: + return sorted(p.name for p in samples_dir().glob("*.pdf")) + + +def guess_schema(path: str | Path) -> str: + truth = Path(path).with_suffix(".truth.json") + if truth.exists(): + return json.loads(truth.read_text())["schema"] + return "contract" if "contract" in Path(path).name.lower() or "agreement" in Path(path).name.lower() else "invoice" + + +def _font(size: int): + for name in ("DejaVuSans-Bold.ttf", "Arial Bold.ttf"): + try: + return ImageFont.truetype(name, size) + except OSError: + continue + return ImageFont.load_default() + + +def draw_document(path: str | Path, boxes: list[tuple[int, tuple, str, str]]) -> Image.Image: + """Render all pages stacked vertically with highlighted boxes. + + ``boxes``: (page, bbox in points, label, kind) where kind is ``field`` or ``cite``. + """ + n_pages = len(pdfium.PdfDocument(str(path))) + pages = [render_page(path, i, RENDER_SCALE).convert("RGBA") for i in range(1, n_pages + 1)] + font = _font(12) + out_pages = [] + for i, img in enumerate(pages, start=1): + overlay = Image.new("RGBA", img.size, (0, 0, 0, 0)) + d = ImageDraw.Draw(overlay) + mine = [(bbox, label, kind) for page, bbox, label, kind in boxes if page == i] + pad = 3 + rects = [tuple(v * RENDER_SCALE + o for v, o in zip(b, (-pad, -pad, pad, pad))) for b, _, _ in mine] + for rect, (_, _, kind) in zip(rects, mine): + if kind == "cite": + d.rectangle(rect, fill=CITE_FILL, outline=CITE_COLOUR + (220,), width=2) + else: + d.rectangle(rect, fill=FIELD_COLOUR + (38,), outline=FIELD_COLOUR + (255,), width=2) + placed: list[tuple[float, float, float, float]] = [] + # outermost boxes first, so a badge stays next to its own box when neighbours compete + order = sorted(zip(rects, mine), key=lambda rm: -rm[0][2] if (rm[0][0] + rm[0][2]) > img.width else rm[0][0]) + for rect, (_, label, kind) in order: + if not label: + continue + colour = (CITE_COLOUR if kind == "cite" else FIELD_COLOUR) + (255,) + bw, bh = d.textlength(label, font=font) + 10, 17 + cy = (rect[1] + rect[3]) / 2 + right = (rect[0] + rect[2]) / 2 > img.width / 2 + # badges sit beside the box, pushed sideways past other boxes and badges so they never cover text + bx = rect[2] + 8 if right else rect[0] - bw - 8 + obstacles = [r for r in rects if r is not rect and r[1] < cy + bh / 2 and r[3] > cy - bh / 2] + placed + moved, shifted = True, False + while moved: + moved = False + for r in obstacles: + if bx < r[2] + 2 and bx + bw > r[0] - 2 and cy - bh / 2 < r[3] and cy + bh / 2 > r[1]: + bx = r[2] + 4 if right else r[0] - bw - 4 + moved = shifted = True + bx = min(max(bx, 2), img.width - bw - 2) + anchor = rect[2] if right else rect[0] + end = bx + bw / 2 + if shifted: # route the leader along the box bottom so it does not cross a neighbour's text + y = rect[3] + 2 + d.line((rect[0] if not right else (rect[0] + rect[2]) / 2, y, end, y), fill=colour, width=1) + d.line((end, y, end, cy + bh / 2), fill=colour, width=1) + else: + d.line((anchor, cy, bx if right else bx + bw, cy), fill=colour, width=1) + d.rounded_rectangle((bx, cy - bh / 2, bx + bw, cy + bh / 2), radius=4, fill=colour) + d.text((bx + 5, cy - bh / 2 + 1), label, font=font, fill=(255, 255, 255, 255)) + placed.append((bx, cy - bh / 2, bx + bw, cy + bh / 2)) + out_pages.append(Image.alpha_composite(img, overlay).convert("RGB")) + width = max(p.width for p in out_pages) + height = sum(p.height for p in out_pages) + PAGE_GAP * (len(out_pages) - 1) + sheet = Image.new("RGB", (width, height), (232, 226, 214)) + y = 0 + for p in out_pages: + sheet.paste(p, (0, y)) + y += p.height + PAGE_GAP + return sheet + + +def field_rows(result: ExtractionResult) -> list[list]: + rows = [] + for i, (name, fv) in enumerate(result.fields.items(), start=1): + value = "" if fv.value is None else (f"{fv.value:,.2f}" if isinstance(fv.value, float) else str(fv.value)) + rows.append([i, name, value, fv.source.page if fv.source else "", fv.source.text if fv.source else "not found"]) + return rows + + +def field_boxes(result: ExtractionResult | None) -> list: + if result is None: + return [] + return [(fv.source.page, fv.source.bbox, str(i), "field") + for i, fv in enumerate(result.fields.values(), start=1) if fv.source] + + +def cite_boxes(answer: Answer | None) -> list: + if answer is None: + return [] + return [(c.page, c.bbox, f"[{c.n}]", "cite") for c in answer.citations] + + +def answer_markdown(answer: Answer) -> str: + lines = [f"**{answer.answer}**", ""] + for c in answer.citations: + snippet = " ".join(c.text.split()) + if len(snippet) > 220: + snippet = snippet[:217] + "..." + lines.append(f"[{c.n}] page {c.page}: {snippet}") + lines.append("") + lines.append(f"Retriever: {answer.retriever} · Model: {answer.model}") + return "\n".join(lines) + + +CSS = """ +:root, .gradio-container { --ms-ink:#1c1a17; --ms-rust:#8a3a24; --ms-paper:#f3efe6; --ms-cream:#faf7f0; } +.gradio-container { background: var(--ms-cream) !important; font-family: Georgia, 'Times New Roman', serif !important; max-width: 1240px !important; } +.ms-header { border-bottom: 1px solid var(--ms-ink); padding: 10px 4px 12px; margin-bottom: 4px; } +.ms-brand { letter-spacing: .22em; font-size: 12px; color: var(--ms-rust); font-weight: 700; } +.ms-header h1 { font-family: Georgia, serif; font-weight: 400; font-size: 30px; margin: 2px 0 2px; color: var(--ms-ink); } +.ms-header p { margin: 0; color: #4a463f; font-style: italic; } +.ms-meta { margin-top: 8px; font-family: ui-monospace, Menlo, monospace; font-size: 12px; color: #6b665e; } +.ms-muted { color: #6b665e; font-size: 12px; } +.ms-card { background: #fff !important; border: 1px solid #ddd5c6 !important; border-radius: 10px !important; } +.ms-section-title { font-size: 12px; letter-spacing: .14em; text-transform: uppercase; color: var(--ms-rust); margin: 2px 0; } +.ms-left { max-height: calc(100vh - 150px); min-height: 560px; overflow-y: auto !important; padding-right: 8px; scroll-behavior: smooth; flex-wrap: nowrap !important; } +.ms-left > * { flex-shrink: 0 !important; } +.ms-table * { font-family: ui-sans-serif, system-ui, -apple-system, 'Segoe UI', sans-serif !important; font-size: 12.5px !important; } +.ms-status, .ms-answer { font-size: 14px; } +.ms-answer strong { color: var(--ms-ink); } +.ms-secondary, button.ms-secondary { background: #fff !important; color: var(--ms-rust) !important; border: 1px solid var(--ms-rust) !important; } +button.primary, .ms-primary button { background: var(--ms-rust) !important; border-color: var(--ms-rust) !important; color: #fff !important; } +footer { display: none !important; } +""" + + +def build_app(settings: Settings | None = None): + import gradio as gr + + settings = settings or Settings.from_env() + cache = ResultCache(settings.cache_dir) + samples = list_samples() + + def pick_sample(name): + if not name: + return gr.update() # dropdown cleared after an upload: keep the uploaded file + return str(samples_dir() / name) + + def pick_upload(file): + return (file if isinstance(file, str) else getattr(file, "name", None)), gr.update(value=None) + + def preview(path): + if not path: + return None, gr.update(), [], "Pick a synthetic sample or upload a PDF.", None, "" + return (draw_document(path, []), guess_schema(path), [], + f"Loaded **{Path(path).name}**. Choose a schema and press **Extract fields**.", None, "") + + def do_extract(path, schema): + if not path: + return [], None, "Pick a sample or upload a PDF first.", None, "" + cached = cache.get(path, "extract", schema, settings) + try: + if cached: + result = ExtractionResult.model_validate(cached) + note = f"Loaded a saved result for this file, computed earlier by `{result.model}`." + else: + result = run_extract(path, schema, settings) + cache.put(path, "extract", schema, settings, result.to_json_dict()) + note = f"Extracted with `{result.model}` ({settings.backend} backend)." + except NoTextLayerError as exc: + return [], draw_document(path, []), f"⚠️ {exc}", None, "" + found = sum(1 for f in result.fields.values() if f.value is not None) + linked = sum(1 for f in result.fields.values() if f.source is not None) + status = f"{note} {found} of {len(result.fields)} fields found, {linked} linked to a source box." + return field_rows(result), draw_document(path, field_boxes(result)), status, result.to_json_dict(), "" + + def do_ask(path, question, result_json): + if not path or not question.strip(): + return "Pick a document and type a question first.", gr.update() + result = ExtractionResult.model_validate(result_json) if result_json else None + cached = cache.get(path, "ask", question.strip(), settings) + try: + if cached: + answer = Answer.model_validate(cached) + else: + answer = run_ask(path, question.strip(), settings) + cache.put(path, "ask", question.strip(), settings, answer.model_dump()) + except NoTextLayerError as exc: + return f"⚠️ {exc}", gr.update() + return answer_markdown(answer), draw_document(path, field_boxes(result) + cite_boxes(answer)) + + with gr.Blocks(title="document-extractor · Modelshifts") as demo: + gr.HTML( + "
MODELSHIFTS

document-extractor

" + "

Turn PDF invoices and contracts into structured fields, each linked to its source.

" + f"
backend: {settings.backend} · model: {settings.model_name} · " + f"v{__version__}
" + ) + current = gr.State(None) + result_state = gr.State(None) + with gr.Row(equal_height=False): + with gr.Column(scale=4, min_width=400, elem_classes="ms-left", elem_id="ms-left"): + gr.HTML("
1 · Document
") + with gr.Row(equal_height=True): + sample = gr.Dropdown(samples, value=None, label="Synthetic sample", interactive=True, scale=3) + upload = gr.UploadButton("Upload PDF", file_types=[".pdf"], type="filepath", scale=1, + elem_classes="ms-secondary") + schema = gr.Radio(["invoice", "contract"], value="invoice", label="Schema") + extract_btn = gr.Button("Extract fields", variant="primary", elem_classes="ms-primary") + status = gr.Markdown("Pick a synthetic sample or upload a PDF.", elem_classes="ms-status") + gr.HTML("
2 · Fields
") + table = gr.Dataframe(headers=["#", "Field", "Value", "Pg", "Source text"], + datatype=["number", "str", "str", "str", "str"], interactive=False, + wrap=True, column_widths=["6%", "27%", "26%", "7%", "34%"], + elem_classes="ms-table") + gr.HTML("
3 · Ask the document
") + with gr.Row(equal_height=True): + question = gr.Textbox(placeholder="e.g. When is payment due?", show_label=False, scale=4, + container=False) + ask_btn = gr.Button("Ask", scale=1, elem_classes="ms-secondary") + answer = gr.Markdown("", elem_classes="ms-answer") + with gr.Column(scale=6, min_width=480): + gr.HTML("
Page view · numbered boxes are field sources, " + "amber boxes are cited passages
") + page = gr.Image(type="pil", show_label=False, height=720, interactive=False, + elem_classes="ms-card", buttons=["fullscreen", "download"]) + gr.HTML("

All sample documents are synthetic. Uploaded files stay on the machine running " + "this demo.

") + + outs = [page, schema, table, status, result_state, answer] + sample.change(pick_sample, [sample], [current]).then(preview, [current], outs) + upload.upload(pick_upload, [upload], [current, sample]).then(preview, [current], outs) + extract_btn.click(do_extract, [current, schema], [table, page, status, result_state, answer]) + ask_btn.click(do_ask, [current, question, result_state], [answer, page]) + question.submit(do_ask, [current, question, result_state], [answer, page]) + return demo + + +def theme(): + import gradio as gr + + return gr.themes.Base( + primary_hue=gr.themes.colors.orange, + neutral_hue=gr.themes.colors.stone, + font=[gr.themes.Font("Georgia"), gr.themes.Font("Times New Roman"), gr.themes.Font("serif")], + font_mono=[gr.themes.Font("ui-monospace"), gr.themes.Font("Menlo"), gr.themes.Font("monospace")], + radius_size=gr.themes.sizes.radius_md, + ) + + +def launch(host: str = "127.0.0.1", port: int = 7860, share: bool = False) -> None: + demo = build_app() + demo.queue().launch(server_name=host, server_port=port, share=share, theme=theme(), css=CSS) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..dae5075 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,35 @@ +import json +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parent.parent +SAMPLES = ROOT / "samples" + + +def truth(name: str) -> dict: + return json.loads((SAMPLES / f"{name}.truth.json").read_text()) + + +class FakeLLM: + """Stands in for a chat model. Returns canned replies and records prompts.""" + + name = "fake-llm" + + def __init__(self, replies): + self.replies = list(replies) if isinstance(replies, (list, tuple)) else [replies] + self.prompts = [] + + def complete(self, messages, max_new_tokens=None): + self.prompts.append(messages) + return self.replies.pop(0) if len(self.replies) > 1 else self.replies[0] + + +@pytest.fixture +def samples() -> Path: + return SAMPLES + + +@pytest.fixture +def fake_llm(): + return FakeLLM diff --git a/tests/test_extract.py b/tests/test_extract.py new file mode 100644 index 0000000..cb85246 --- /dev/null +++ b/tests/test_extract.py @@ -0,0 +1,66 @@ +import json + +from document_extractor.extract import extract +from document_extractor.grounding import ground +from document_extractor.pdf import load_pdf +from document_extractor.schema import load_schema + +from conftest import truth + + +def verbatim_reply(name): + """What a well-behaved model returns: values copied as printed on the page.""" + t = truth(name)["fields"] + return { + "invoice-northwind-classic": { + "invoice_number": "NWP-2026-0412", "invoice_date": "14 March 2026", "due_date": "13 April 2026", + "seller_name": "Northwind Parts Pte. Ltd.", "buyer_name": "Example Logistics Pte. Ltd.", + "currency": "SGD", "subtotal": f"{t['subtotal']:,.2f}", "tax": f"{t['tax']:,.2f}", + "total": f"{t['total']:,.2f}", + }, + }[name] + + +def test_extract_types_and_sources(samples, fake_llm): + name = "invoice-northwind-classic" + llm = fake_llm("Here you go:\n```json\n" + json.dumps(verbatim_reply(name)) + "\n```") + doc = load_pdf(samples / f"{name}.pdf") + result = extract(doc, load_schema("invoice"), llm, backend="test") + want = truth(name)["fields"] + assert result.values() == want + for field, fv in result.fields.items(): + assert fv.source is not None, field + assert fv.source.page == 1 + assert result.fields["invoice_number"].source.text == "NWP-2026-0412" + assert result.fields["total"].source.text.endswith(f"{want['total']:,.2f}") + # the prompt carries the document text and the field list + prompt = llm.prompts[0][1]["content"] + assert "NWP-2026-0412" in prompt and "invoice_number (string)" in prompt + + +def test_normalised_values_are_still_grounded(samples, fake_llm): + name = "invoice-demo-freight-compact" + want = truth(name)["fields"] + reply = dict(want) # model returned ISO dates and plain floats instead of the printed text + result = extract(load_pdf(samples / f"{name}.pdf"), load_schema("invoice"), fake_llm(json.dumps(reply))) + assert result.values() == want + assert result.fields["invoice_date"].source.text == "03.02.2026" + assert result.fields["total"].source is not None + + +def test_missing_and_bad_json(samples, fake_llm): + doc = load_pdf(samples / "contract-supply.pdf") + result = extract(doc, load_schema("contract"), fake_llm("sorry, no idea")) + assert all(fv.value is None and fv.source is None for fv in result.fields.values()) + + +def test_hallucinated_value_has_no_source(samples): + doc = load_pdf(samples / "invoice-northwind-classic.pdf") + assert ground(doc, "ZZ-9999-XYZ", "ZZ-9999-XYZ") is None + + +def test_repeated_value_prefers_labelled_copy(samples): + doc = load_pdf(samples / "contract-master-services.pdf") + src = ground(doc, "Singapore", "Singapore", "string", hint="governing_law The jurisdiction whose laws govern") + lines = [w for w in doc.words if w.page == src.page and abs(w.bbox[1] - src.bbox[1]) < 2] + assert "governed" in " ".join(w.text for w in lines) diff --git a/tests/test_ocr_and_web.py b/tests/test_ocr_and_web.py new file mode 100644 index 0000000..40240ad --- /dev/null +++ b/tests/test_ocr_and_web.py @@ -0,0 +1,24 @@ +from document_extractor.ocr import parse_vision_lines, split_line + + +def test_split_line_shares_box(): + words = split_line("Total SGD 5,689.11", (100, 10, 280, 22), page=1) + assert [w.text for w in words] == ["Total", "SGD", "5,689.11"] + assert words[0].bbox[0] == 100 and abs(words[-1].bbox[2] - 280) < 1e-6 + assert all(w.bbox[1] == 10 and w.bbox[3] == 22 for w in words) + + +def test_parse_vision_lines_scales_to_points(): + content = 'Sure: [{"text": "Invoice No. AES-5120", "box": [100, 200, 400, 220]}]' + words = parse_vision_lines(content, size=(1190, 1684), page=1, scale=2.0) + assert words[0].text == "Invoice" and words[-1].text == "AES-5120" + assert abs(words[0].bbox[0] - 59.5) < 0.01 # 100/1000 * 595 pt + + +def test_draw_document_marks_boxes(samples): + from document_extractor.webapp import draw_document + + plain = draw_document(samples / "invoice-placeholder-two-page.pdf", []) + marked = draw_document(samples / "invoice-placeholder-two-page.pdf", [(2, (300, 300, 400, 320), "1", "field")]) + assert plain.size == marked.size and plain.height > plain.width # two pages stacked + assert plain.tobytes() != marked.tobytes() diff --git a/tests/test_pdf.py b/tests/test_pdf.py new file mode 100644 index 0000000..c5b81d9 --- /dev/null +++ b/tests/test_pdf.py @@ -0,0 +1,47 @@ +import pytest + +from document_extractor.pdf import NoTextLayerError, document_text, load_pdf, render_page + + +def test_words_have_page_and_bbox(samples): + doc = load_pdf(samples / "invoice-placeholder-two-page.pdf") + assert [p.number for p in doc.pages] == [1, 2] + for page in doc.pages: + assert page.words + for w in page.words: + x0, top, x1, bottom = w.bbox + assert w.page == page.number + assert 0 <= x0 < x1 <= page.width + assert 0 <= top < bottom <= page.height + + +def test_document_text_has_page_markers(samples): + text = document_text(load_pdf(samples / "invoice-placeholder-two-page.pdf")) + assert "[Page 1]" in text and "[Page 2]" in text + assert "PS-INV-30027" in text + + +def test_scanned_pdf_without_ocr_gives_clear_message(samples): + with pytest.raises(NoTextLayerError) as err: + load_pdf(samples / "invoice-acme-scanned.pdf") + assert "MS_VISION_MODEL" in str(err.value) and "MS_BACKEND=hf" in str(err.value) + + +def test_scanned_pdf_uses_ocr_engine(samples): + from document_extractor.models import Word + + class FakeOCR: + name = "fake-ocr" + + def recognise(self, image, page, scale): + assert image.width > 500 + return [Word(text="AES-5120", page=page, bbox=(100, 100, 160, 112))] + + doc = load_pdf(samples / "invoice-acme-scanned.pdf", ocr=FakeOCR()) + assert doc.pages[0].ocr is True + assert doc.pages[0].words[0].text == "AES-5120" + + +def test_render_page(samples): + img = render_page(samples / "invoice-northwind-classic.pdf", 1, scale=1.0) + assert abs(img.width - 595) <= 1 and abs(img.height - 842) <= 1 diff --git a/tests/test_qa_cli.py b/tests/test_qa_cli.py new file mode 100644 index 0000000..181392f --- /dev/null +++ b/tests/test_qa_cli.py @@ -0,0 +1,103 @@ +import csv +import json + +from document_extractor import cli +from document_extractor.pdf import load_pdf +from document_extractor.qa import ask +from document_extractor.retrieval import BM25Retriever, chunk + +from conftest import FakeLLM, truth + + +def test_chunks_cover_pages(samples): + doc = load_pdf(samples / "contract-master-services.pdf") + passages = chunk(doc) + assert len(passages) > 3 + assert any("sixty (60) days" in p.text for p in passages) + + +def test_ask_cites_retrieved_passage(samples): + doc = load_pdf(samples / "contract-master-services.pdf") + llm = FakeLLM("The notice period starts on the day the notice is received [1].") + answer = ask(doc, "When does the notice period start?", llm, BM25Retriever()) + assert answer.citations and answer.citations[0].n == 1 + assert "notice" in answer.citations[0].text.lower() + assert answer.retriever == "bm25" + assert "Passages:" in llm.prompts[0][1]["content"] + + +def test_ask_not_found_has_no_citation(samples): + doc = load_pdf(samples / "contract-supply.pdf") + answer = ask(doc, "What colour is the logo?", FakeLLM("Not found in the document."), BM25Retriever()) + assert answer.citations == [] + stray = ask(doc, "What colour is the logo?", FakeLLM("Not found in the document. [2]"), BM25Retriever()) + assert stray.answer == "Not found in the document." and stray.citations == [] + + +def _patch_llm(monkeypatch): + from document_extractor import pipeline + + def fake_get_llm(settings): + def reply_for(messages): + content = messages[1]["content"] + for t in ("invoice-northwind-classic", "invoice-example-logistics-banner"): + fields = truth(t)["fields"] + if fields["invoice_number"] in content: + return json.dumps(fields) + return "{}" + + class L: + name = "fake-llm" + + def complete(self, messages, max_new_tokens=None): + return reply_for(messages) + + return L() + + monkeypatch.setattr(pipeline, "get_llm", fake_get_llm) + monkeypatch.setenv("MS_BACKEND", "api") + monkeypatch.setenv("MS_API_BASE", "http://localhost:0/v1") + monkeypatch.setenv("MS_MODEL", "fake") + + +def test_cli_extract(samples, monkeypatch, capsys): + _patch_llm(monkeypatch) + assert cli.main(["extract", str(samples / "invoice-northwind-classic.pdf"), "--schema", "invoice"]) == 0 + out = json.loads(capsys.readouterr().out) + assert out["schema"] == "invoice" + assert out["fields"]["invoice_number"]["source"]["page"] == 1 + + +def test_cli_batch_csv_and_json(samples, monkeypatch, tmp_path): + _patch_llm(monkeypatch) + folder = tmp_path / "in" + folder.mkdir() + for n in ("invoice-northwind-classic", "invoice-example-logistics-banner", "invoice-acme-scanned"): + (folder / f"{n}.pdf").write_bytes((samples / f"{n}.pdf").read_bytes()) + out_csv = tmp_path / "results.csv" + assert cli.main(["batch", str(folder), "--schema", "invoice", "--out", str(out_csv)]) == 0 + rows = list(csv.DictReader(out_csv.open())) + assert len(rows) == 3 + by_file = {r["file"]: r for r in rows} + assert by_file["invoice-northwind-classic.pdf"]["invoice_number"] == "NWP-2026-0412" + assert "no text layer" in by_file["invoice-acme-scanned.pdf"]["error"] # api backend without a vision model + out_json = tmp_path / "results.json" + assert cli.main(["batch", str(folder), "--schema", "invoice", "--out", str(out_json)]) == 0 + assert len(json.loads(out_json.read_text())) == 2 + + +def test_cli_bad_backend(monkeypatch, capsys, samples): + monkeypatch.setenv("MS_BACKEND", "nope") + assert cli.main(["extract", str(samples / "contract-supply.pdf"), "--schema", "contract"]) == 2 + assert "MS_BACKEND" in capsys.readouterr().err + + +def test_wrong_citation_number_is_repaired(samples): + doc = load_pdf(samples / "contract-master-services.pdf") + retriever = BM25Retriever() + passages = retriever.search("When does the notice period start?", chunk(doc), 4) + wrong = next(i for i, p in enumerate(passages, 1) if "received" not in p.text) + llm = FakeLLM(f"The notice period starts on the day the notice is received [{wrong}].") + answer = ask(doc, "When does the notice period start?", llm, retriever) + assert "received" in answer.citations[0].text + assert f"[{answer.citations[0].n}]" in answer.answer diff --git a/tests/test_schema.py b/tests/test_schema.py new file mode 100644 index 0000000..9e341a7 --- /dev/null +++ b/tests/test_schema.py @@ -0,0 +1,43 @@ +import pytest + +from document_extractor.schema import coerce, load_schema, parse_date, parse_integer, parse_number + + +def test_bundled_schemas_load(): + inv = load_schema("invoice") + con = load_schema("schemas/contract.yaml") + assert "total" in inv.field_names and inv.name == "invoice" + assert "notice_period_days" in con.field_names + + +@pytest.mark.parametrize("text,want", [ + ("14 March 2026", "2026-03-14"), + ("2026-05-02", "2026-05-02"), + ("03.02.2026", "2026-02-03"), + ("20/01/2026", "2026-01-20"), + ("March 14, 2026", "2026-03-14"), + ("1st July 2026", "2026-07-01"), + ("soon", None), +]) +def test_parse_date(text, want): + assert parse_date(text) == want + + +@pytest.mark.parametrize("text,want", [ + ("5,496.68", 5496.68), ("SGD 453.85", 453.85), ("1.234,56", 1234.56), ("12,50", 12.5), ("n/a", None), +]) +def test_parse_number(text, want): + assert parse_number(text) == want + + +@pytest.mark.parametrize("text,want", [ + ("sixty (60) days", 60), ("24 months", 24), ("twelve months", 12), ("45", 45), +]) +def test_parse_integer(text, want): + assert parse_integer(text) == want + + +def test_coerce_nulls(): + assert coerce(None, "string") is None + assert coerce("null", "number") is None + assert coerce(12, "number") == 12.0