DX fixes: markdown cleanup, error propagation, extract_batch, dataclass

README.md

@@ -6,13 +6,13 @@
 </picture>
 
 [![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/)
 [![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
 
 </div>
 
-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) </br>
-**🎯 Accurate** — matches SOTA quality: 0.862–0.873 ROUGE-5 F1 on WebMainBench </br>
-**🪶 Small** — the recommended model is 210M params, fits on any GPU </br>
-**💸 Cheap** — clean 1 billion pages for ~$7,900 vs ~$159,000 for the leading decoder </br>
-**📦 Simple** — `pip install pulpie`, then `Extractor().extract(html)` </br>
-**🔌 Batched** — overlapped CPU+GPU pipeline scales across multiple GPUs </br>
+- **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):
 
@@ -153,5 +153,5 @@ If you use Pulpie in your research, please cite:
 ---
 
 <div align="center">
-Built by <a href="https://github.com/chonkie-inc">Chonkie</a>, the open-source work behind <a href="https://usefeyn.com">Feyn</a>.
+Built by <a href="https://usefeyn.com">Feyn</a>.
 </div>

src/pulpie/extractor.py

@@ -2,10 +2,14 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
+
 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,
     load_model_and_tokenizer,
     predictions_to_labels,
@@ -33,27 +37,29 @@ def __init__(
         model: str = DEFAULT_MODEL,
         device: str | None = None,
         max_tokens: int = 8192,
+        max_batch_tokens: int = 16384,
     ):
         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
+        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."""
         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,
         )
 
@@ -64,71 +70,131 @@ 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,
         )
 
-    @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,
             bos_token_id=self.tokenizer.bos_token_id,
             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)
 
-    def _to_markdown(self, html: str) -> str:
-        """Convert HTML to markdown."""
-        try:
-            import html2text
+    @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
 
-            h = html2text.HTML2Text(bodywidth=0)
-            h.ignore_links = False
-            h.ignore_images = False
-            return h.handle(html).strip()
-        except ImportError:
-            return html
+        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:

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
+
+# <img> 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"<img\b[^>]*?"
+    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 `<img>` 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()

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]: