From b1d55ad0b78798feef312a1ea794725c722560c2 Mon Sep 17 00:00:00 2001 From: "Bhavnick @ Chonkie" Date: Mon, 29 Jun 2026 21:54:24 -0700 Subject: [PATCH 1/6] README footer: built by Feyn --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 67d9d80..4f3b8f8 100644 --- a/README.md +++ b/README.md @@ -153,5 +153,5 @@ If you use Pulpie in your research, please cite: ---
-Built by Chonkie, the open-source work behind Feyn. +Built by Feyn.
From 7817f99ee36ee7cff8005ba339061cce385f26c5 Mon Sep 17 00:00:00 2001 From: "Bhavnick @ Chonkie" Date: Mon, 29 Jun 2026 21:56:08 -0700 Subject: [PATCH 2/6] README: static 'python 3.9+' badge instead of per-version list --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 4f3b8f8..73b1957 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ [![PyPI version](https://img.shields.io/pypi/v/pulpie.svg)](https://pypi.org/project/pulpie/) -[![Python versions](https://img.shields.io/pypi/pyversions/pulpie.svg)](https://pypi.org/project/pulpie/) +[![Python](https://img.shields.io/badge/python-3.9%2B-blue.svg)](https://pypi.org/project/pulpie/) [![License](https://img.shields.io/github/license/chonkie-inc/pulpie.svg)](https://github.com/chonkie-inc/pulpie/blob/main/LICENSE) [![Downloads](https://static.pepy.tech/badge/pulpie)](https://pepy.tech/project/pulpie) [![Blog](https://img.shields.io/badge/blog-read%20the%20writeup-E34C26.svg)](https://usefeyn.com/blog/pulpie-pareto-optimal-models-for-cleaning-the-web/) From 6093e7468fe48618ec0f41c7d3fee43df2d7d9bb Mon Sep 17 00:00:00 2001 From: "Bhavnick @ Chonkie" Date: Mon, 29 Jun 2026 21:59:06 -0700 Subject: [PATCH 3/6] README: formal bullets, qualify 20x as L4-specific, drop em dashes - Convert emoji feature list to plain markdown bullets - Qualify the 20x faster/cheaper claim as 'on an L4 GPU' (it's 20.1x on L4, 7.1x on A100 per the blog) so it doesn't read as universal - Remove all em/en dashes from prose per style preference --- README.md | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 73b1957..f53c94d 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ [![Blog](https://img.shields.io/badge/blog-read%20the%20writeup-E34C26.svg)](https://usefeyn.com/blog/pulpie-pareto-optimal-models-for-cleaning-the-web/) [![GitHub stars](https://img.shields.io/github/stars/chonkie-inc/pulpie.svg)](https://github.com/chonkie-inc/pulpie/stargazers) -_Pareto-optimal models for cleaning the web — extract main content from HTML at one twentieth the cost._ +_Pareto-optimal models for cleaning the web. Extract main content from HTML at one twentieth the cost._ [Install](#installation) • [Usage](#usage) • @@ -23,14 +23,14 @@ _Pareto-optimal models for cleaning the web — extract main content from HTML a -Pulpie extracts the main content from raw HTML — stripping navigation, ads, sidebars, and footers — using small encoder models that label every block in a single forward pass. It approaches state-of-the-art extraction quality while running up to **20x faster** and **20x cheaper** than autoregressive extractors. +Pulpie extracts the main content from raw HTML, stripping navigation, ads, sidebars, and footers. It uses small encoder models that label every block in a single forward pass, approaching state-of-the-art extraction quality while running up to 20x faster and 20x cheaper than autoregressive extractors on an L4 GPU. -**⚡ Fast** — an encoder labels every block in one forward pass (13.7 pages/sec on an L4)
-**🎯 Accurate** — matches SOTA quality: 0.862–0.873 ROUGE-5 F1 on WebMainBench
-**🪶 Small** — the recommended model is 210M params, fits on any GPU
-**💸 Cheap** — clean 1 billion pages for ~$7,900 vs ~$159,000 for the leading decoder
-**📦 Simple** — `pip install pulpie`, then `Extractor().extract(html)`
-**🔌 Batched** — overlapped CPU+GPU pipeline scales across multiple GPUs
+- **Fast.** An encoder labels every block in one forward pass (13.7 pages/sec on an L4). +- **Accurate.** Matches state-of-the-art quality: 0.862 to 0.873 ROUGE-5 F1 on WebMainBench. +- **Small.** The recommended model is 210M parameters and fits on any GPU. +- **Cheap.** Clean 1 billion pages for ~$7,900, versus ~$159,000 for the leading decoder. +- **Simple.** Run `pip install pulpie`, then `Extractor().extract(html)`. +- **Batched.** An overlapped CPU and GPU pipeline scales across multiple GPUs. ## Installation @@ -94,7 +94,7 @@ All three models are built on [EuroBERT](https://arxiv.org/abs/2503.05500), shar | Model | Hugging Face | Params | ROUGE-5 F1 | Notes | |-------|--------------|--------|------------|-------| -| **Orange Small** | [`chonkie-ai/pulpie-orange-small`](https://huggingface.co/chonkie-ai/pulpie-orange-small) | 210M | 0.862 | **Recommended** — best size-to-quality ratio | +| **Orange Small** | [`chonkie-ai/pulpie-orange-small`](https://huggingface.co/chonkie-ai/pulpie-orange-small) | 210M | 0.862 | **Recommended**, best size-to-quality ratio | | Orange Base | [`chonkie-ai/pulpie-orange-base`](https://huggingface.co/chonkie-ai/pulpie-orange-base) | 610M | 0.863 | Distilled from Large | | Orange Large | [`chonkie-ai/pulpie-orange-large`](https://huggingface.co/chonkie-ai/pulpie-orange-large) | 2.1B | 0.873 | Teacher (highest quality) | @@ -104,12 +104,12 @@ All three models are built on [EuroBERT](https://arxiv.org/abs/2503.05500), shar Pulpie keeps the "read the page" approach of model-based extractors but moves the bottleneck from memory bandwidth to compute by using an encoder instead of a decoder. The pipeline runs in four stages: -1. **Simplify** — remove scripts, styles, and formatting noise; tag each content block with a unique ID. -2. **Chunk** — split, tokenize, and pack blocks into chunks of up to 8,192 tokens (≈80% of pages fit in one chunk). -3. **Classify** — a single encoder forward pass labels every block as content or boilerplate. -4. **Reconstruct** — return the kept blocks as HTML, or convert them to Markdown. +1. **Simplify.** Remove scripts, styles, and formatting noise; tag each content block with a unique ID. +2. **Chunk.** Split, tokenize, and pack blocks into chunks of up to 8,192 tokens (≈80% of pages fit in one chunk). +3. **Classify.** A single encoder forward pass labels every block as content or boilerplate. +4. **Reconstruct.** Return the kept blocks as HTML, or convert them to Markdown. -A decoder emits labels one token at a time, re-reading the full model from GPU memory each step. An encoder runs one dense forward pass over the whole input — so the gap widens on bandwidth-limited GPUs (7x faster than Dripper on A100, 20x on L4). +A decoder emits labels one token at a time, re-reading the full model from GPU memory each step. An encoder runs one dense forward pass over the whole input, so the gap widens on bandwidth-limited GPUs (7x faster than Dripper on A100, 20x on L4). ## Benchmarks @@ -121,8 +121,8 @@ Quality on the English subset of [WebMainBench](https://github.com/opendatalab/W | Dripper | 0.6B | 0.864 | 135 | | **Pulpie Orange Base** | 610M | 0.863 | 36 | | **Pulpie Orange Small** | 210M | 0.862 | 45 | -| magic-html | — | 0.700 | 384 | -| Trafilatura | — | 0.619 | 16 | +| magic-html | - | 0.700 | 384 | +| Trafilatura | - | 0.619 | 16 | Speed and cost (Pulpie Orange Small vs Dripper, 1 billion pages): From b2c6bd3d872a3f0dd042b1e8c8f7b93ffad94212 Mon Sep 17 00:00:00 2001 From: "Bhavnick @ Chonkie" Date: Mon, 29 Jun 2026 23:10:51 -0700 Subject: [PATCH 4/6] Auto-detect Apple MPS device Add default_device() (cuda -> mps -> cpu) and use it in Extractor and Pipeline when device is unspecified. Previously device=None fell back straight to CPU on Macs even when MPS was available, silently leaving Apple acceleration unused. --- src/pulpie/extractor.py | 3 ++- src/pulpie/model_utils.py | 9 +++++++++ src/pulpie/pipeline.py | 3 ++- 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/pulpie/extractor.py b/src/pulpie/extractor.py index 25fcfba..b8d60b2 100644 --- a/src/pulpie/extractor.py +++ b/src/pulpie/extractor.py @@ -6,6 +6,7 @@ from pulpie.chunker import extract_blocks, pack_chunks, tokenize_blocks from pulpie.model_utils import ( + default_device, extract_item_ids, load_model_and_tokenizer, predictions_to_labels, @@ -35,7 +36,7 @@ def __init__( max_tokens: int = 8192, ): if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" + device = default_device() self.device = torch.device(device) self.max_tokens = max_tokens diff --git a/src/pulpie/model_utils.py b/src/pulpie/model_utils.py index fb0a1fb..1d275cd 100644 --- a/src/pulpie/model_utils.py +++ b/src/pulpie/model_utils.py @@ -23,6 +23,15 @@ def resolve_model_id(model: str) -> str: return MODELS.get(model, model) +def default_device() -> str: + """Pick the best available device: CUDA, then Apple MPS, then CPU.""" + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return "cpu" + + def load_model_and_tokenizer( model_id: str, device: torch.device ) -> tuple[torch.nn.Module, AutoTokenizer, int]: diff --git a/src/pulpie/pipeline.py b/src/pulpie/pipeline.py index a557f1c..deaff99 100644 --- a/src/pulpie/pipeline.py +++ b/src/pulpie/pipeline.py @@ -20,6 +20,7 @@ from pulpie.chunker import SEP_TOKEN, extract_blocks, pack_chunks, tokenize_blocks from pulpie.model_utils import ( + default_device, extract_item_ids, load_model_and_tokenizer, predictions_to_labels, @@ -208,7 +209,7 @@ def __init__( if torch.cuda.is_available(): devices = [f"cuda:{i}" for i in range(torch.cuda.device_count())] else: - devices = ["cpu"] + devices = [default_device()] # "mps" on Apple, else "cpu" elif isinstance(devices, str): devices = [devices] self.devices = [torch.device(d) for d in devices] From 76e74a32dcc40eb0dfd9eea8cc56238d892dce32 Mon Sep 17 00:00:00 2001 From: "Bhavnick @ Chonkie" Date: Mon, 29 Jun 2026 23:20:22 -0700 Subject: [PATCH 5/6] DX fixes: markdown cleanup, error propagation, max_batch_tokens, page_id docs - Add pulpie/markdown.py with shared to_markdown() that strips tracking-pixel / spacer tags (1x1, trans/blank/spacer, data-URI) before html2text. Replaces the duplicated html2text block in extractor + pipeline. - Pipeline now propagates per-page errors to PageResult.error (simplify failures and dropped pages were previously indistinguishable from empty pages). - Wire Pipeline(max_batch_tokens=...) through to inference (was silently hardcoded to 16384 in _infer_and_push). - Document page_id as an opaque echo (ordering is positional) and that error is set on failure. --- src/pulpie/extractor.py | 19 +++----------- src/pulpie/markdown.py | 45 ++++++++++++++++++++++++++++++++ src/pulpie/pipeline.py | 57 +++++++++++++++++++++++++++-------------- 3 files changed, 86 insertions(+), 35 deletions(-) create mode 100644 src/pulpie/markdown.py diff --git a/src/pulpie/extractor.py b/src/pulpie/extractor.py index b8d60b2..0a60cc0 100644 --- a/src/pulpie/extractor.py +++ b/src/pulpie/extractor.py @@ -5,6 +5,7 @@ import torch from pulpie.chunker import extract_blocks, pack_chunks, tokenize_blocks +from pulpie.markdown import to_markdown from pulpie.model_utils import ( default_device, extract_item_ids, @@ -50,11 +51,10 @@ def extract(self, html: str) -> ExtractionResult: simplified, map_html = simplify(html) labels = self._classify(simplified) main_html = extract_main_html(map_html, labels) - markdown = self._to_markdown(main_html) return ExtractionResult( html=main_html, - markdown=markdown, + markdown=to_markdown(main_html), labels=labels, ) @@ -65,11 +65,10 @@ def extract_from_simplified( labels = self._classify(simplified_html) source = map_html if map_html is not None else simplified_html main_html = extract_main_html(source, labels) - markdown = self._to_markdown(main_html) return ExtractionResult( html=main_html, - markdown=markdown, + markdown=to_markdown(main_html), labels=labels, ) @@ -108,18 +107,6 @@ def _classify(self, simplified_html: str) -> dict[str, str]: return predictions_to_labels(item_ids, predictions) - def _to_markdown(self, html: str) -> str: - """Convert HTML to markdown.""" - try: - import html2text - - h = html2text.HTML2Text(bodywidth=0) - h.ignore_links = False - h.ignore_images = False - return h.handle(html).strip() - except ImportError: - return html - class ExtractionResult: """Result of content extraction.""" diff --git a/src/pulpie/markdown.py b/src/pulpie/markdown.py new file mode 100644 index 0000000..ea57e27 --- /dev/null +++ b/src/pulpie/markdown.py @@ -0,0 +1,45 @@ +"""HTML -> Markdown conversion shared by Extractor and Pipeline. + +Strips tracking-pixel / spacer images before conversion so they don't leak into +the output as empty `![](...)` tags, then renders with html2text. If html2text +isn't installed (it's an optional `[markdown]` extra), the cleaned HTML is +returned unchanged. +""" + +from __future__ import annotations + +import re + +# tags that carry no real content: 1x1 spacers, blank/transparent gifs, +# or data-URI placeholders. Matched case-insensitively on the whole tag. +_SPACER_IMG = re.compile( + r"]*?" + r'(?:(?:width|height)\s*=\s*["\']?\s*1\b' # width=1 / height=1 + r'|src\s*=\s*["\'][^"\']*' + r"(?:trans(?:parent)?|spacer|blank|pixel|1x1|clear)[^\"']*[\"']" + r"|src\s*=\s*[\"']data:image[^\"']*[\"'])" + r"[^>]*>", + re.IGNORECASE, +) + + +def strip_spacer_images(html: str) -> str: + """Remove zero-content / tracking-pixel `` tags.""" + return _SPACER_IMG.sub("", html) + + +def to_markdown(html: str) -> str: + """Convert main-content HTML to Markdown. + + Returns the (spacer-stripped) HTML unchanged if html2text isn't available. + """ + cleaned = strip_spacer_images(html) + try: + import html2text + except ImportError: + return cleaned + + h = html2text.HTML2Text(bodywidth=0) + h.ignore_links = False + h.ignore_images = False + return h.handle(cleaned).strip() diff --git a/src/pulpie/pipeline.py b/src/pulpie/pipeline.py index deaff99..5bce45b 100644 --- a/src/pulpie/pipeline.py +++ b/src/pulpie/pipeline.py @@ -19,6 +19,7 @@ import torch from pulpie.chunker import SEP_TOKEN, extract_blocks, pack_chunks, tokenize_blocks +from pulpie.markdown import to_markdown from pulpie.model_utils import ( default_device, extract_item_ids, @@ -32,7 +33,14 @@ @dataclass class PageInput: - """Input to the pipeline.""" + """Input to the pipeline. + + ``page_id`` is an opaque identifier echoed back on the matching + ``PageResult`` so you can correlate inputs and outputs. It is *not* used for + ordering (results are returned in input order regardless). If you don't set + it, it defaults to 0 and every result will report ``page_id=0`` — set it + (e.g. the input's list index) if you need to tell results apart by id. + """ html: str page_id: int = 0 @@ -41,7 +49,12 @@ class PageInput: @dataclass class PageResult: - """Output from the pipeline.""" + """Output from the pipeline. + + Returned in the same order as the input list. ``error`` is set (and the + other fields left empty) when a page failed to process; it is ``None`` on + success. + """ page_id: int labels: dict[str, str] @@ -135,26 +148,22 @@ def _cpu_prepare( def _postprocess( - page_id: int, batch_idx: int, labels: dict[str, str], map_html: str + page_id: int, + batch_idx: int, + labels: dict[str, str], + map_html: str, + error: str | None = None, ) -> tuple[int, PageResult]: """CPU post-processing: reconstruct HTML + convert to markdown.""" main_html = extract_main_html(map_html, labels) - - try: - import html2text - - h = html2text.HTML2Text(bodywidth=0) - h.ignore_links = False - h.ignore_images = False - markdown = h.handle(main_html).strip() - except ImportError: - markdown = main_html + markdown = to_markdown(main_html) return batch_idx, PageResult( page_id=page_id, labels=labels, html=main_html, markdown=markdown, + error=error, ) @@ -236,6 +245,7 @@ def __init__( sep_token_id=sep_token_id, device=device, pad_id=model_inst.config.pad_token_id or 0, + max_batch_tokens=self.max_batch_tokens, ) ) @@ -288,7 +298,13 @@ def extract_batch(self, pages: list[PageInput]) -> list[PageResult]: # Fill gaps for i in range(n_pages): if results[i] is None: - results[i] = PageResult(page_id=pages[i].page_id, labels={}, html="", markdown="") + results[i] = PageResult( + page_id=pages[i].page_id, + labels={}, + html="", + markdown="", + error="no result produced (page dropped before post-processing)", + ) return results # type: ignore[return-value] @@ -335,7 +351,9 @@ def _stage2_gpu_worker( _infer_and_push(state, pending, out_queue) return if item.error or not item.chunks: - out_queue.put((item.page_id, item.batch_idx, {}, item.map_html)) + out_queue.put( + (item.page_id, item.batch_idx, {}, item.map_html, item.error) + ) else: pending.append(item) except queue.Empty: @@ -363,8 +381,8 @@ def _stage3_postprocess(self, in_queue: queue.Queue, results: list[PageResult | if item is None: continue - page_id, batch_idx, labels, map_html = item - future = pool.submit(_postprocess, page_id, batch_idx, labels, map_html) + page_id, batch_idx, labels, map_html, error = item + future = pool.submit(_postprocess, page_id, batch_idx, labels, map_html, error) pending_futures[future] = batch_idx _collect_futures(pending_futures, results) @@ -382,6 +400,7 @@ class _GPUWorkerState: sep_token_id: int device: torch.device pad_id: int + max_batch_tokens: int def _infer_and_push( @@ -390,7 +409,7 @@ def _infer_and_push( out_queue: queue.Queue, ) -> None: """Run batched inference on accumulated pages and push results.""" - max_batch_tokens = 16384 + max_batch_tokens = state.max_batch_tokens # Flatten chunks, sort by length all_chunks: list[tuple[list[int], list[int], int]] = [] @@ -435,7 +454,7 @@ def _infer_and_push( if idx < len(preds): predictions[block_idx] = preds[idx] labels = predictions_to_labels(page.item_ids, predictions) - out_queue.put((page.page_id, page.batch_idx, labels, page.map_html)) + out_queue.put((page.page_id, page.batch_idx, labels, page.map_html, None)) def _collect_futures(pending: dict, results: list) -> None: From a36d8e44f923cc34b3b4a94ae7aa34bac9284c06 Mon Sep 17 00:00:00 2001 From: "Bhavnick @ Chonkie" Date: Mon, 29 Jun 2026 23:43:45 -0700 Subject: [PATCH 6/6] Add Extractor.extract_batch + make ExtractionResult a dataclass - extract_batch(htmls): convenience list-in/list-out API that packs chunks across pages into shared forward passes. Verified on GPU box: 0/13 label mismatches vs per-page extract. Documented as a convenience (not faster than a loop with this model's eager attention; use Pipeline for scale throughput). - _classify stays simple/sequential (memory-safe); shared _chunk helper. - ExtractionResult is now a @dataclass (consistent with PageInput/PageResult), keeping the compact custom __repr__. Verified on gym-pro6000 (CUDA): correct across all 13 fixtures; 65 parity tests pass. --- src/pulpie/extractor.py | 122 ++++++++++++++++++++++++++++++++-------- 1 file changed, 100 insertions(+), 22 deletions(-) diff --git a/src/pulpie/extractor.py b/src/pulpie/extractor.py index 0a60cc0..f549fc4 100644 --- a/src/pulpie/extractor.py +++ b/src/pulpie/extractor.py @@ -2,6 +2,8 @@ from __future__ import annotations +from dataclasses import dataclass + import torch from pulpie.chunker import extract_blocks, pack_chunks, tokenize_blocks @@ -35,16 +37,19 @@ def __init__( model: str = DEFAULT_MODEL, device: str | None = None, max_tokens: int = 8192, + max_batch_tokens: int = 16384, ): if device is None: device = default_device() self.device = torch.device(device) self.max_tokens = max_tokens + self.max_batch_tokens = max_batch_tokens model_id = resolve_model_id(model) self.model, self.tokenizer, self.sep_token_id = load_model_and_tokenizer( model_id, self.device ) + self.pad_id = self.model.config.pad_token_id or 0 def extract(self, html: str) -> ExtractionResult: """Extract main content from raw HTML.""" @@ -72,16 +77,43 @@ def extract_from_simplified( labels=labels, ) - @torch.no_grad() - def _classify(self, simplified_html: str) -> dict[str, str]: - """Classify each block as main/other.""" - blocks = extract_blocks(simplified_html) - if not blocks: - return {} - - item_ids = extract_item_ids(blocks) + def extract_batch(self, htmls: list[str]) -> list[ExtractionResult]: + """Extract many pages with one call. Returns results in input order. + + A convenience wrapper over :meth:`extract` that packs chunks from + different pages into shared forward passes. Note: with this model's eager + attention it is not meaningfully faster than calling :meth:`extract` in a + loop (a single pass already saturates the GPU); it's offered for + ergonomics. For large-scale, multi-GPU throughput use + :class:`pulpie.Pipeline`. + """ + prepared = [] # (item_ids, n_blocks, map_html, [(chunk_ids, block_indices), ...]) + all_chunks = [] # (page_idx, chunk_ids, block_indices) + for page_idx, html in enumerate(htmls): + simplified, map_html = simplify(html) + blocks = extract_blocks(simplified) + item_ids = extract_item_ids(blocks) + chunks = self._chunk(blocks) + prepared.append((item_ids, len(blocks), map_html)) + for chunk_ids, block_indices in chunks: + all_chunks.append((page_idx, chunk_ids, block_indices)) + + predictions = [[0] * n for _, n, _ in prepared] + self._run_batched(all_chunks, predictions) + + results = [] + for (item_ids, _n, map_html), preds in zip(prepared, predictions): + labels = predictions_to_labels(item_ids, preds) + main_html = extract_main_html(map_html, labels) + results.append( + ExtractionResult(html=main_html, markdown=to_markdown(main_html), labels=labels) + ) + return results + + def _chunk(self, blocks: list[str]) -> list[tuple[list[int], list[int]]]: + """Tokenize blocks and pack them into model-sized chunks.""" block_token_ids = tokenize_blocks(blocks, self.tokenizer) - chunks = pack_chunks( + return pack_chunks( block_token_ids, max_tokens=self.max_tokens, sep_token_id=self.sep_token_id, @@ -89,34 +121,80 @@ def _classify(self, simplified_html: str) -> dict[str, str]: eos_token_id=self.tokenizer.eos_token_id, ) - predictions = [0] * len(blocks) + @torch.no_grad() + def _classify(self, simplified_html: str) -> dict[str, str]: + """Classify each block of a single page as main/other (sequential).""" + blocks = extract_blocks(simplified_html) + if not blocks: + return {} - for chunk_ids, block_indices in chunks: + item_ids = extract_item_ids(blocks) + predictions = [0] * len(blocks) + for chunk_ids, block_indices in self._chunk(blocks): input_ids = torch.tensor([chunk_ids], dtype=torch.long, device=self.device) attention_mask = torch.ones_like(input_ids) - - outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) - logits = outputs.logits[0] - + logits = self.model(input_ids=input_ids, attention_mask=attention_mask).logits[0] sep_positions = (input_ids[0] == self.sep_token_id).nonzero(as_tuple=True)[0] preds = logits[sep_positions].argmax(dim=-1).cpu().tolist() - for i, block_idx in enumerate(block_indices): if i < len(preds): predictions[block_idx] = preds[i] return predictions_to_labels(item_ids, predictions) + @torch.no_grad() + def _run_batched(self, all_chunks, predictions) -> None: + """Run chunks (tagged with their page index) in length-sorted, padded batches. + + Memory for eager attention scales as ``batch * max_len^2``, so the batch + size is capped by the squared longest length to stay safe on long chunks + while grouping short ones (the common case across many single-chunk pages). + """ + if not all_chunks: + return + budget = self.max_batch_tokens * self.max_batch_tokens + ordered = sorted(all_chunks, key=lambda c: len(c[1])) + + batch: list = [] + for item in ordered: + cand = batch + [item] + max_len = len(cand[-1][1]) # ascending -> last is longest + if batch and len(cand) * max_len * max_len > budget: + self._infer_batch(batch, predictions) + batch = [item] + else: + batch = cand + if batch: + self._infer_batch(batch, predictions) + + def _infer_batch(self, batch, predictions) -> None: + """Run one padded batch of (page_idx, chunk_ids, block_indices); write preds.""" + max_len = max(len(c[1]) for c in batch) + input_ids = torch.full( + (len(batch), max_len), self.pad_id, dtype=torch.long, device=self.device + ) + attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long, device=self.device) + for row, (_page_idx, chunk_ids, _bi) in enumerate(batch): + input_ids[row, : len(chunk_ids)] = torch.tensor(chunk_ids, dtype=torch.long) + attention_mask[row, : len(chunk_ids)] = 1 + + logits = self.model(input_ids=input_ids, attention_mask=attention_mask).logits + for row, (page_idx, _chunk_ids, block_indices) in enumerate(batch): + sep_positions = (input_ids[row] == self.sep_token_id).nonzero(as_tuple=True)[0] + preds = logits[row][sep_positions].argmax(dim=-1).cpu().tolist() + for j, block_idx in enumerate(block_indices): + if j < len(preds): + predictions[page_idx][block_idx] = preds[j] + + +@dataclass class ExtractionResult: """Result of content extraction.""" - __slots__ = ("html", "labels", "markdown") - - def __init__(self, html: str, markdown: str, labels: dict[str, str]): - self.html = html - self.markdown = markdown - self.labels = labels + html: str + markdown: str + labels: dict[str, str] @property def n_blocks(self) -> int: