prettypipeline-ocr 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- prettypipeline_ocr-0.1.0/LICENSE +21 -0
- prettypipeline_ocr-0.1.0/PKG-INFO +99 -0
- prettypipeline_ocr-0.1.0/README.md +71 -0
- prettypipeline_ocr-0.1.0/pyproject.toml +41 -0
- prettypipeline_ocr-0.1.0/setup.cfg +4 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline/__init__.py +1 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline/cli.py +65 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline/extract.py +138 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline/ocr.py +132 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline_ocr.egg-info/PKG-INFO +99 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline_ocr.egg-info/SOURCES.txt +14 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline_ocr.egg-info/dependency_links.txt +1 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline_ocr.egg-info/entry_points.txt +2 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline_ocr.egg-info/requires.txt +13 -0
- prettypipeline_ocr-0.1.0/src/prettypipeline_ocr.egg-info/top_level.txt +1 -0
- prettypipeline_ocr-0.1.0/tests/test_review.py +21 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 PrettyPipeline contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: prettypipeline-ocr
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Turn PDFs into structured JSON with local OCR and a cheap cloud LLM.
|
|
5
|
+
Author: greatshoor
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/virajshoor/PrettyPipeline
|
|
8
|
+
Keywords: ocr,pdf,llm,extraction,json
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Topic :: Scientific/Engineering :: Image Recognition
|
|
11
|
+
Requires-Python: <3.14,>=3.10
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
License-File: LICENSE
|
|
14
|
+
Requires-Dist: pymupdf>=1.24
|
|
15
|
+
Requires-Dist: transformers<5,>=4.57.1
|
|
16
|
+
Requires-Dist: torch
|
|
17
|
+
Requires-Dist: torchvision
|
|
18
|
+
Requires-Dist: pillow
|
|
19
|
+
Requires-Dist: openai>=1.0
|
|
20
|
+
Requires-Dist: einops
|
|
21
|
+
Requires-Dist: addict
|
|
22
|
+
Requires-Dist: easydict
|
|
23
|
+
Requires-Dist: safetensors
|
|
24
|
+
Requires-Dist: tqdm
|
|
25
|
+
Requires-Dist: matplotlib
|
|
26
|
+
Requires-Dist: requests
|
|
27
|
+
Dynamic: license-file
|
|
28
|
+
|
|
29
|
+
# PrettyPipeline
|
|
30
|
+
|
|
31
|
+
Turn PDFs into structured JSON. OCR runs **locally**; field extraction uses a **cheap cloud LLM**.
|
|
32
|
+
|
|
33
|
+
## Why this is cheap
|
|
34
|
+
|
|
35
|
+
Cloud OCR APIs charge per page. PrettyPipeline does the expensive vision work on your GPU with [baidu/Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) (MIT, 3B params, 32K context — a whole PDF in one `infer_multi()` pass). The only API call is [GPT-5.4 nano](https://developers.openai.com/api/docs/models/gpt-5.4-nano) on the resulting text (~$0.20 / 1M input tokens), prompted with **your** JSON schema.
|
|
36
|
+
|
|
37
|
+
You pay for a small text completion, not for pixels.
|
|
38
|
+
|
|
39
|
+
```
|
|
40
|
+
PDF → PyMuPDF pages → Unlimited-OCR (local GPU) → raw text
|
|
41
|
+
↓
|
|
42
|
+
GPT-5.4 nano (OpenAI) → JSON
|
|
43
|
+
↓
|
|
44
|
+
validation → needs_review flags
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Nulls, model-marked uncertain fields, and OCR that looks garbled are flagged for a human instead of being silently accepted.
|
|
48
|
+
|
|
49
|
+
## Install
|
|
50
|
+
|
|
51
|
+
Python **3.10–3.13** (3.12 recommended). NVIDIA GPU (CUDA) or Apple Silicon (MPS). CPU works but is slow.
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
python3.12 -m venv .venv
|
|
55
|
+
source .venv/bin/activate
|
|
56
|
+
pip install prettypipeline-ocr
|
|
57
|
+
# or from this repo:
|
|
58
|
+
pip install -e .
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
On NVIDIA Linux, install a CUDA build of PyTorch first if the default wheel is CPU-only:
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
|
|
65
|
+
pip install -e .
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
The first run downloads `baidu/Unlimited-OCR` from Hugging Face (~6GB).
|
|
69
|
+
|
|
70
|
+
## Usage
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
export OPENAI_API_KEY=sk-...
|
|
74
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json
|
|
75
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json -o out.json
|
|
76
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --ocr-only
|
|
77
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --device mps
|
|
78
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --max-length 2048
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
`--schema` is any JSON Schema. Swap the file to extract a different document type — nothing is hardcoded.
|
|
82
|
+
|
|
83
|
+
Device is auto-detected (`cuda` → `mps` → `cpu`). Override with `--device`.
|
|
84
|
+
|
|
85
|
+
On Apple Silicon, Unlimited-OCR can start looping (it hardcodes CUDA). Use `--max-length 2048` for short docs. NVIDIA CUDA is the model's native path.
|
|
86
|
+
|
|
87
|
+
## Output
|
|
88
|
+
|
|
89
|
+
```json
|
|
90
|
+
{
|
|
91
|
+
"data": { "date": "…", "vendor": "…", "line_items": [], "total": 0 },
|
|
92
|
+
"needs_review": [{ "field": "tax", "reason": "null" }],
|
|
93
|
+
"ocr_text": "…"
|
|
94
|
+
}
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
## License
|
|
98
|
+
|
|
99
|
+
MIT. Unlimited-OCR is also MIT.
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
# PrettyPipeline
|
|
2
|
+
|
|
3
|
+
Turn PDFs into structured JSON. OCR runs **locally**; field extraction uses a **cheap cloud LLM**.
|
|
4
|
+
|
|
5
|
+
## Why this is cheap
|
|
6
|
+
|
|
7
|
+
Cloud OCR APIs charge per page. PrettyPipeline does the expensive vision work on your GPU with [baidu/Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) (MIT, 3B params, 32K context — a whole PDF in one `infer_multi()` pass). The only API call is [GPT-5.4 nano](https://developers.openai.com/api/docs/models/gpt-5.4-nano) on the resulting text (~$0.20 / 1M input tokens), prompted with **your** JSON schema.
|
|
8
|
+
|
|
9
|
+
You pay for a small text completion, not for pixels.
|
|
10
|
+
|
|
11
|
+
```
|
|
12
|
+
PDF → PyMuPDF pages → Unlimited-OCR (local GPU) → raw text
|
|
13
|
+
↓
|
|
14
|
+
GPT-5.4 nano (OpenAI) → JSON
|
|
15
|
+
↓
|
|
16
|
+
validation → needs_review flags
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
Nulls, model-marked uncertain fields, and OCR that looks garbled are flagged for a human instead of being silently accepted.
|
|
20
|
+
|
|
21
|
+
## Install
|
|
22
|
+
|
|
23
|
+
Python **3.10–3.13** (3.12 recommended). NVIDIA GPU (CUDA) or Apple Silicon (MPS). CPU works but is slow.
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
python3.12 -m venv .venv
|
|
27
|
+
source .venv/bin/activate
|
|
28
|
+
pip install prettypipeline-ocr
|
|
29
|
+
# or from this repo:
|
|
30
|
+
pip install -e .
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
On NVIDIA Linux, install a CUDA build of PyTorch first if the default wheel is CPU-only:
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
|
|
37
|
+
pip install -e .
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
The first run downloads `baidu/Unlimited-OCR` from Hugging Face (~6GB).
|
|
41
|
+
|
|
42
|
+
## Usage
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
export OPENAI_API_KEY=sk-...
|
|
46
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json
|
|
47
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json -o out.json
|
|
48
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --ocr-only
|
|
49
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --device mps
|
|
50
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --max-length 2048
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
`--schema` is any JSON Schema. Swap the file to extract a different document type — nothing is hardcoded.
|
|
54
|
+
|
|
55
|
+
Device is auto-detected (`cuda` → `mps` → `cpu`). Override with `--device`.
|
|
56
|
+
|
|
57
|
+
On Apple Silicon, Unlimited-OCR can start looping (it hardcodes CUDA). Use `--max-length 2048` for short docs. NVIDIA CUDA is the model's native path.
|
|
58
|
+
|
|
59
|
+
## Output
|
|
60
|
+
|
|
61
|
+
```json
|
|
62
|
+
{
|
|
63
|
+
"data": { "date": "…", "vendor": "…", "line_items": [], "total": 0 },
|
|
64
|
+
"needs_review": [{ "field": "tax", "reason": "null" }],
|
|
65
|
+
"ocr_text": "…"
|
|
66
|
+
}
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
## License
|
|
70
|
+
|
|
71
|
+
MIT. Unlimited-OCR is also MIT.
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "prettypipeline-ocr"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Turn PDFs into structured JSON with local OCR and a cheap cloud LLM."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
requires-python = ">=3.10,<3.14"
|
|
12
|
+
authors = [{ name = "greatshoor" }]
|
|
13
|
+
keywords = ["ocr", "pdf", "llm", "extraction", "json"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Topic :: Scientific/Engineering :: Image Recognition",
|
|
17
|
+
]
|
|
18
|
+
dependencies = [
|
|
19
|
+
"pymupdf>=1.24",
|
|
20
|
+
"transformers>=4.57.1,<5",
|
|
21
|
+
"torch",
|
|
22
|
+
"torchvision",
|
|
23
|
+
"pillow",
|
|
24
|
+
"openai>=1.0",
|
|
25
|
+
"einops",
|
|
26
|
+
"addict",
|
|
27
|
+
"easydict",
|
|
28
|
+
"safetensors",
|
|
29
|
+
"tqdm",
|
|
30
|
+
"matplotlib",
|
|
31
|
+
"requests",
|
|
32
|
+
]
|
|
33
|
+
|
|
34
|
+
[project.urls]
|
|
35
|
+
Homepage = "https://github.com/virajshoor/PrettyPipeline"
|
|
36
|
+
|
|
37
|
+
[project.scripts]
|
|
38
|
+
prettypipeline = "prettypipeline.cli:main"
|
|
39
|
+
|
|
40
|
+
[tool.setuptools.packages.find]
|
|
41
|
+
where = ["src"]
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.1.0"
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
"""CLI: prettypipeline run <file.pdf> --schema <schema.json>"""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import json
|
|
7
|
+
import sys
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
from prettypipeline.extract import needs_review, require_api_key, structure
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def main(argv: list[str] | None = None) -> int:
|
|
14
|
+
p = argparse.ArgumentParser(
|
|
15
|
+
prog="prettypipeline",
|
|
16
|
+
description="PDF → local OCR → cheap LLM JSON extraction.",
|
|
17
|
+
)
|
|
18
|
+
sub = p.add_subparsers(dest="cmd", required=True)
|
|
19
|
+
run = sub.add_parser("run", help="OCR a PDF and extract JSON using a schema")
|
|
20
|
+
run.add_argument("pdf", type=Path)
|
|
21
|
+
run.add_argument("--schema", type=Path, required=True)
|
|
22
|
+
run.add_argument("-o", "--output", type=Path, help="Write JSON here (also printed)")
|
|
23
|
+
run.add_argument("--ocr-only", action="store_true", help="Skip the cloud structuring step")
|
|
24
|
+
run.add_argument("--device", choices=("cuda", "mps", "cpu"), default="", help="Override auto device")
|
|
25
|
+
run.add_argument("--dpi", type=int, default=300)
|
|
26
|
+
run.add_argument("--max-length", type=int, default=32768, help="OCR generation cap (model supports 32768)")
|
|
27
|
+
args = p.parse_args(argv)
|
|
28
|
+
|
|
29
|
+
if args.cmd != "run":
|
|
30
|
+
p.print_help()
|
|
31
|
+
return 2
|
|
32
|
+
if not args.pdf.is_file():
|
|
33
|
+
print(f"not a file: {args.pdf}", file=sys.stderr)
|
|
34
|
+
return 2
|
|
35
|
+
schema = json.loads(args.schema.read_text())
|
|
36
|
+
from prettypipeline.ocr import ocr_pdf, pick_device
|
|
37
|
+
|
|
38
|
+
device = args.device or str(pick_device())
|
|
39
|
+
print(f"OCR device: {device}", file=sys.stderr)
|
|
40
|
+
text = ocr_pdf(str(args.pdf), dpi=args.dpi, device=args.device, max_length=args.max_length)
|
|
41
|
+
if args.ocr_only:
|
|
42
|
+
result = {"ocr_text": text, "data": None, "needs_review": []}
|
|
43
|
+
_emit(result, args.output)
|
|
44
|
+
return 0
|
|
45
|
+
require_api_key()
|
|
46
|
+
extracted = structure(text, schema)
|
|
47
|
+
result = {
|
|
48
|
+
"data": extracted["data"],
|
|
49
|
+
"needs_review": needs_review(extracted["data"], text, extracted["uncertain_fields"]),
|
|
50
|
+
"ocr_text": text,
|
|
51
|
+
}
|
|
52
|
+
_emit(result, args.output)
|
|
53
|
+
return 0
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _emit(result: dict, output: Path | None) -> None:
|
|
57
|
+
blob = json.dumps(result, indent=2, ensure_ascii=False)
|
|
58
|
+
print(blob)
|
|
59
|
+
if output:
|
|
60
|
+
output.parent.mkdir(parents=True, exist_ok=True)
|
|
61
|
+
output.write_text(blob + "\n")
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
if __name__ == "__main__":
|
|
65
|
+
raise SystemExit(main())
|
|
@@ -0,0 +1,138 @@
|
|
|
1
|
+
"""Cheap cloud structuring (GPT-5.4 nano) plus local review flags."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import os
|
|
7
|
+
import re
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
from openai import OpenAI
|
|
11
|
+
|
|
12
|
+
MODEL = "gpt-5.4-nano"
|
|
13
|
+
API_KEY_ENV = "OPENAI_API_KEY"
|
|
14
|
+
|
|
15
|
+
_GARBLED = re.compile(r"[\ufffd]|[^\w\s.,:$€£¥%/@+#()\\-]{4,}")
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def require_api_key() -> str:
|
|
19
|
+
key = os.environ.get(API_KEY_ENV, "").strip()
|
|
20
|
+
if not key:
|
|
21
|
+
raise SystemExit(
|
|
22
|
+
f"{API_KEY_ENV} is not set. Export it, then re-run:\n"
|
|
23
|
+
f" export {API_KEY_ENV}=sk-...\n"
|
|
24
|
+
f" prettypipeline run <file.pdf> --schema <schema.json>"
|
|
25
|
+
)
|
|
26
|
+
return key
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _clean_ocr(ocr_text: str) -> str:
|
|
30
|
+
text = re.sub(r"<\|det\|>.*?<\|/det\|>", " ", ocr_text, flags=re.DOTALL)
|
|
31
|
+
return re.sub(r"<PAGE>", "\n", text)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def structure(ocr_text: str, schema: dict[str, Any], api_key: str | None = None) -> dict[str, Any]:
|
|
35
|
+
client = OpenAI(api_key=api_key or require_api_key())
|
|
36
|
+
messages = [
|
|
37
|
+
{
|
|
38
|
+
"role": "system",
|
|
39
|
+
"content": (
|
|
40
|
+
"Extract fields from OCR text into JSON. "
|
|
41
|
+
"Return exactly: {\"data\": <object matching the schema>, "
|
|
42
|
+
"\"uncertain_fields\": [<dotted paths you are not confident about>]}. "
|
|
43
|
+
"Use null when a required or optional field is missing or unreadable. "
|
|
44
|
+
"Do not invent values. Prefer null over a guess. "
|
|
45
|
+
"Ignore repeated garbage tokens from OCR degeneration."
|
|
46
|
+
),
|
|
47
|
+
},
|
|
48
|
+
{
|
|
49
|
+
"role": "user",
|
|
50
|
+
"content": (
|
|
51
|
+
"JSON schema (target shape for data):\n"
|
|
52
|
+
f"{json.dumps(schema, indent=2)}\n\n"
|
|
53
|
+
"OCR text:\n"
|
|
54
|
+
f"{_clean_ocr(ocr_text)}"
|
|
55
|
+
),
|
|
56
|
+
},
|
|
57
|
+
]
|
|
58
|
+
resp = client.chat.completions.create(
|
|
59
|
+
model=MODEL,
|
|
60
|
+
messages=messages,
|
|
61
|
+
response_format={"type": "json_object"},
|
|
62
|
+
)
|
|
63
|
+
raw = resp.choices[0].message.content or "{}"
|
|
64
|
+
parsed = json.loads(raw)
|
|
65
|
+
data = parsed.get("data", parsed)
|
|
66
|
+
uncertain = parsed.get("uncertain_fields") or []
|
|
67
|
+
if not isinstance(uncertain, list):
|
|
68
|
+
uncertain = []
|
|
69
|
+
cleaned_uncertain = []
|
|
70
|
+
for x in uncertain:
|
|
71
|
+
path = str(x)
|
|
72
|
+
if path.startswith("data."):
|
|
73
|
+
path = path[5:]
|
|
74
|
+
cleaned_uncertain.append(path)
|
|
75
|
+
return {"data": data, "uncertain_fields": cleaned_uncertain}
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def _walk(obj: Any, prefix: str = "") -> list[tuple[str, Any]]:
|
|
79
|
+
out: list[tuple[str, Any]] = []
|
|
80
|
+
if isinstance(obj, dict):
|
|
81
|
+
for k, v in obj.items():
|
|
82
|
+
path = f"{prefix}.{k}" if prefix else k
|
|
83
|
+
out.extend(_walk(v, path))
|
|
84
|
+
elif isinstance(obj, list):
|
|
85
|
+
for i, v in enumerate(obj):
|
|
86
|
+
out.extend(_walk(v, f"{prefix}[{i}]"))
|
|
87
|
+
else:
|
|
88
|
+
out.append((prefix, obj))
|
|
89
|
+
return out
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def looks_garbled(value: str) -> bool:
|
|
93
|
+
if not value or not str(value).strip():
|
|
94
|
+
return False
|
|
95
|
+
s = str(value)
|
|
96
|
+
if s.count("\ufffd") >= 1:
|
|
97
|
+
return True
|
|
98
|
+
if _GARBLED.search(s):
|
|
99
|
+
return True
|
|
100
|
+
alnum = sum(c.isalnum() for c in s)
|
|
101
|
+
return len(s) >= 8 and alnum / len(s) < 0.35
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def ocr_near_field_garbled(ocr_text: str, value: str) -> bool:
|
|
105
|
+
if not value or not ocr_text:
|
|
106
|
+
return False
|
|
107
|
+
needle = str(value).strip()
|
|
108
|
+
if len(needle) < 3:
|
|
109
|
+
return False
|
|
110
|
+
idx = ocr_text.lower().find(needle.lower())
|
|
111
|
+
if idx < 0:
|
|
112
|
+
return False
|
|
113
|
+
window = ocr_text[max(0, idx - 40) : idx + len(needle) + 40]
|
|
114
|
+
return looks_garbled(window)
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def needs_review(
|
|
118
|
+
data: Any,
|
|
119
|
+
ocr_text: str = "",
|
|
120
|
+
uncertain: list[str] | None = None,
|
|
121
|
+
) -> list[dict[str, str]]:
|
|
122
|
+
flags: list[dict[str, str]] = []
|
|
123
|
+
seen: set[tuple[str, str]] = set()
|
|
124
|
+
|
|
125
|
+
def add(field: str, reason: str) -> None:
|
|
126
|
+
key = (field, reason)
|
|
127
|
+
if key not in seen:
|
|
128
|
+
seen.add(key)
|
|
129
|
+
flags.append({"field": field, "reason": reason})
|
|
130
|
+
|
|
131
|
+
for path in uncertain or []:
|
|
132
|
+
add(path, "uncertain")
|
|
133
|
+
for path, value in _walk(data):
|
|
134
|
+
if value is None:
|
|
135
|
+
add(path, "null")
|
|
136
|
+
elif isinstance(value, str) and (looks_garbled(value) or ocr_near_field_garbled(ocr_text, value)):
|
|
137
|
+
add(path, "garbled_ocr")
|
|
138
|
+
return flags
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
"""Local OCR via baidu/Unlimited-OCR (Transformers). CUDA, Apple Silicon MPS, or CPU."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
import shutil
|
|
7
|
+
import tempfile
|
|
8
|
+
from functools import lru_cache
|
|
9
|
+
|
|
10
|
+
import pymupdf as fitz
|
|
11
|
+
import torch
|
|
12
|
+
from transformers import AutoModel, AutoTokenizer
|
|
13
|
+
|
|
14
|
+
MODEL_ID = "baidu/Unlimited-OCR"
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def pick_device(explicit: str | None = None) -> torch.device:
|
|
18
|
+
if explicit:
|
|
19
|
+
return torch.device(explicit)
|
|
20
|
+
if torch.cuda.is_available():
|
|
21
|
+
return torch.device("cuda")
|
|
22
|
+
if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
|
|
23
|
+
return torch.device("mps")
|
|
24
|
+
return torch.device("cpu")
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def _patch_cuda_calls(device: torch.device) -> None:
|
|
28
|
+
"""Unlimited-OCR hardcodes Tensor.cuda() and autocast('cuda')."""
|
|
29
|
+
if device.type == "cuda":
|
|
30
|
+
return
|
|
31
|
+
|
|
32
|
+
target = device
|
|
33
|
+
|
|
34
|
+
def _tensor_cuda(self, *args, **kwargs):
|
|
35
|
+
return self.to(target)
|
|
36
|
+
|
|
37
|
+
def _module_cuda(self, *args, **kwargs):
|
|
38
|
+
return self.to(target)
|
|
39
|
+
|
|
40
|
+
_orig_autocast = torch.autocast
|
|
41
|
+
|
|
42
|
+
def _autocast(device_type=None, dtype=None, *args, **kwargs):
|
|
43
|
+
if device_type == "cuda":
|
|
44
|
+
if device.type == "mps":
|
|
45
|
+
kwargs.pop("device_type", None)
|
|
46
|
+
try:
|
|
47
|
+
return _orig_autocast("mps", dtype=dtype, *args, **kwargs)
|
|
48
|
+
except TypeError:
|
|
49
|
+
return _orig_autocast("cpu", dtype=dtype, *args, **kwargs)
|
|
50
|
+
return _orig_autocast("cpu", dtype=dtype, *args, **kwargs)
|
|
51
|
+
return _orig_autocast(device_type, dtype=dtype, *args, **kwargs)
|
|
52
|
+
|
|
53
|
+
torch.Tensor.cuda = _tensor_cuda # type: ignore[method-assign]
|
|
54
|
+
torch.nn.Module.cuda = _module_cuda # type: ignore[method-assign]
|
|
55
|
+
torch.autocast = _autocast # type: ignore[assignment]
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _dtype_for(device: torch.device) -> torch.dtype:
|
|
59
|
+
if device.type == "cuda":
|
|
60
|
+
return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
|
61
|
+
if device.type == "mps":
|
|
62
|
+
return torch.bfloat16
|
|
63
|
+
return torch.float32
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def pdf_to_images(pdf_path: str, dpi: int = 300) -> tuple[list[str], str]:
|
|
67
|
+
"""Rasterize PDF pages the way Unlimited-OCR documents: PyMuPDF at `dpi`."""
|
|
68
|
+
tmp_dir = tempfile.mkdtemp(prefix="pdf_ocr_")
|
|
69
|
+
doc = fitz.open(pdf_path)
|
|
70
|
+
mat = fitz.Matrix(dpi / 72, dpi / 72)
|
|
71
|
+
paths = []
|
|
72
|
+
try:
|
|
73
|
+
for i, page in enumerate(doc):
|
|
74
|
+
out = os.path.join(tmp_dir, f"page_{i + 1:04d}.png")
|
|
75
|
+
page.get_pixmap(matrix=mat).save(out)
|
|
76
|
+
paths.append(out)
|
|
77
|
+
except Exception:
|
|
78
|
+
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
79
|
+
raise
|
|
80
|
+
finally:
|
|
81
|
+
doc.close()
|
|
82
|
+
if not paths:
|
|
83
|
+
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
84
|
+
raise ValueError(f"no pages in {pdf_path}")
|
|
85
|
+
return paths, tmp_dir
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
@lru_cache(maxsize=1)
|
|
89
|
+
def load_model(device_str: str = "") -> tuple[object, object, torch.device]:
|
|
90
|
+
device = pick_device(device_str or None)
|
|
91
|
+
_patch_cuda_calls(device)
|
|
92
|
+
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
|
|
93
|
+
dtype = _dtype_for(device)
|
|
94
|
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
|
95
|
+
model = AutoModel.from_pretrained(
|
|
96
|
+
MODEL_ID,
|
|
97
|
+
trust_remote_code=True,
|
|
98
|
+
use_safetensors=True,
|
|
99
|
+
torch_dtype=dtype,
|
|
100
|
+
)
|
|
101
|
+
model = model.eval().to(device)
|
|
102
|
+
return model, tokenizer, device
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def ocr_pdf(
|
|
106
|
+
pdf_path: str,
|
|
107
|
+
dpi: int = 300,
|
|
108
|
+
device: str = "",
|
|
109
|
+
output_dir: str | None = None,
|
|
110
|
+
max_length: int = 32768,
|
|
111
|
+
) -> str:
|
|
112
|
+
model, tokenizer, _ = load_model(device)
|
|
113
|
+
paths, tmp_dir = pdf_to_images(pdf_path, dpi=dpi)
|
|
114
|
+
out = output_dir or tempfile.mkdtemp(prefix="ocr_out_")
|
|
115
|
+
own_out = output_dir is None
|
|
116
|
+
try:
|
|
117
|
+
text, _tokens = model.infer_multi(
|
|
118
|
+
tokenizer,
|
|
119
|
+
prompt="<image>Multi page parsing.",
|
|
120
|
+
image_files=paths,
|
|
121
|
+
output_path=out,
|
|
122
|
+
image_size=1024,
|
|
123
|
+
max_length=max_length,
|
|
124
|
+
no_repeat_ngram_size=35,
|
|
125
|
+
ngram_window=1024,
|
|
126
|
+
save_results=False,
|
|
127
|
+
)
|
|
128
|
+
return text
|
|
129
|
+
finally:
|
|
130
|
+
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
131
|
+
if own_out:
|
|
132
|
+
shutil.rmtree(out, ignore_errors=True)
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: prettypipeline-ocr
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Turn PDFs into structured JSON with local OCR and a cheap cloud LLM.
|
|
5
|
+
Author: greatshoor
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/virajshoor/PrettyPipeline
|
|
8
|
+
Keywords: ocr,pdf,llm,extraction,json
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Topic :: Scientific/Engineering :: Image Recognition
|
|
11
|
+
Requires-Python: <3.14,>=3.10
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
License-File: LICENSE
|
|
14
|
+
Requires-Dist: pymupdf>=1.24
|
|
15
|
+
Requires-Dist: transformers<5,>=4.57.1
|
|
16
|
+
Requires-Dist: torch
|
|
17
|
+
Requires-Dist: torchvision
|
|
18
|
+
Requires-Dist: pillow
|
|
19
|
+
Requires-Dist: openai>=1.0
|
|
20
|
+
Requires-Dist: einops
|
|
21
|
+
Requires-Dist: addict
|
|
22
|
+
Requires-Dist: easydict
|
|
23
|
+
Requires-Dist: safetensors
|
|
24
|
+
Requires-Dist: tqdm
|
|
25
|
+
Requires-Dist: matplotlib
|
|
26
|
+
Requires-Dist: requests
|
|
27
|
+
Dynamic: license-file
|
|
28
|
+
|
|
29
|
+
# PrettyPipeline
|
|
30
|
+
|
|
31
|
+
Turn PDFs into structured JSON. OCR runs **locally**; field extraction uses a **cheap cloud LLM**.
|
|
32
|
+
|
|
33
|
+
## Why this is cheap
|
|
34
|
+
|
|
35
|
+
Cloud OCR APIs charge per page. PrettyPipeline does the expensive vision work on your GPU with [baidu/Unlimited-OCR](https://huggingface.co/baidu/Unlimited-OCR) (MIT, 3B params, 32K context — a whole PDF in one `infer_multi()` pass). The only API call is [GPT-5.4 nano](https://developers.openai.com/api/docs/models/gpt-5.4-nano) on the resulting text (~$0.20 / 1M input tokens), prompted with **your** JSON schema.
|
|
36
|
+
|
|
37
|
+
You pay for a small text completion, not for pixels.
|
|
38
|
+
|
|
39
|
+
```
|
|
40
|
+
PDF → PyMuPDF pages → Unlimited-OCR (local GPU) → raw text
|
|
41
|
+
↓
|
|
42
|
+
GPT-5.4 nano (OpenAI) → JSON
|
|
43
|
+
↓
|
|
44
|
+
validation → needs_review flags
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Nulls, model-marked uncertain fields, and OCR that looks garbled are flagged for a human instead of being silently accepted.
|
|
48
|
+
|
|
49
|
+
## Install
|
|
50
|
+
|
|
51
|
+
Python **3.10–3.13** (3.12 recommended). NVIDIA GPU (CUDA) or Apple Silicon (MPS). CPU works but is slow.
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
python3.12 -m venv .venv
|
|
55
|
+
source .venv/bin/activate
|
|
56
|
+
pip install prettypipeline-ocr
|
|
57
|
+
# or from this repo:
|
|
58
|
+
pip install -e .
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
On NVIDIA Linux, install a CUDA build of PyTorch first if the default wheel is CPU-only:
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
|
|
65
|
+
pip install -e .
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
The first run downloads `baidu/Unlimited-OCR` from Hugging Face (~6GB).
|
|
69
|
+
|
|
70
|
+
## Usage
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
export OPENAI_API_KEY=sk-...
|
|
74
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json
|
|
75
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json -o out.json
|
|
76
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --ocr-only
|
|
77
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --device mps
|
|
78
|
+
prettypipeline run invoice.pdf --schema examples/invoice.schema.json --max-length 2048
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
`--schema` is any JSON Schema. Swap the file to extract a different document type — nothing is hardcoded.
|
|
82
|
+
|
|
83
|
+
Device is auto-detected (`cuda` → `mps` → `cpu`). Override with `--device`.
|
|
84
|
+
|
|
85
|
+
On Apple Silicon, Unlimited-OCR can start looping (it hardcodes CUDA). Use `--max-length 2048` for short docs. NVIDIA CUDA is the model's native path.
|
|
86
|
+
|
|
87
|
+
## Output
|
|
88
|
+
|
|
89
|
+
```json
|
|
90
|
+
{
|
|
91
|
+
"data": { "date": "…", "vendor": "…", "line_items": [], "total": 0 },
|
|
92
|
+
"needs_review": [{ "field": "tax", "reason": "null" }],
|
|
93
|
+
"ocr_text": "…"
|
|
94
|
+
}
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
## License
|
|
98
|
+
|
|
99
|
+
MIT. Unlimited-OCR is also MIT.
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
src/prettypipeline/__init__.py
|
|
5
|
+
src/prettypipeline/cli.py
|
|
6
|
+
src/prettypipeline/extract.py
|
|
7
|
+
src/prettypipeline/ocr.py
|
|
8
|
+
src/prettypipeline_ocr.egg-info/PKG-INFO
|
|
9
|
+
src/prettypipeline_ocr.egg-info/SOURCES.txt
|
|
10
|
+
src/prettypipeline_ocr.egg-info/dependency_links.txt
|
|
11
|
+
src/prettypipeline_ocr.egg-info/entry_points.txt
|
|
12
|
+
src/prettypipeline_ocr.egg-info/requires.txt
|
|
13
|
+
src/prettypipeline_ocr.egg-info/top_level.txt
|
|
14
|
+
tests/test_review.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
prettypipeline
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
from prettypipeline.extract import looks_garbled, needs_review
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def test_null_uncertain_garbled():
|
|
5
|
+
data = {"date": None, "vendor": "Acme", "total": "@@@@####~~~~xx"}
|
|
6
|
+
flags = needs_review(data, ocr_text="Acme invoice", uncertain=["vendor"])
|
|
7
|
+
reasons = {(f["field"], f["reason"]) for f in flags}
|
|
8
|
+
assert ("date", "null") in reasons
|
|
9
|
+
assert ("vendor", "uncertain") in reasons
|
|
10
|
+
assert ("total", "garbled_ocr") in reasons
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def test_clean_string_not_garbled():
|
|
14
|
+
assert not looks_garbled("Acme Corp")
|
|
15
|
+
assert looks_garbled("bad\ufffdtext")
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
if __name__ == "__main__":
|
|
19
|
+
test_null_uncertain_garbled()
|
|
20
|
+
test_clean_string_not_garbled()
|
|
21
|
+
print("ok")
|