htmlsift 0.1.0__py3-none-any.whl

htmlsift/__init__.py

@@ -0,0 +1,9 @@
+"""htmlsift -- web main-content extraction (HTML in, clean text/HTML/markdown out).
+
+    from htmlsift import Extractor, extract
+"""
+from .extractor import Extractor, extract
+
+__version__ = "0.1.0"
+
+__all__ = ["Extractor", "extract"]

htmlsift/_artifacts.py

@@ -0,0 +1,64 @@
+"""Where the model bundles live and how the package fetches them.
+
+Bundles are published to a public HuggingFace repo and pinned by commit SHA per
+release, so a given htmlsift version always fetches byte-identical files. Fetching
+uses hf_hub_download, which caches to the standard HF cache; a local directory can
+override the fetch during development.
+"""
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+
+REPO = "koivualeksi/htmlsift-release"
+REVISION = "73ac2b978aef356e9e5772640a0c9444f2b67705"   # release commit with both mini/
+                         # and base/ bundles. Pin per release so a version fetches byte-
+                         # identical files; None (dev only) would track the default branch.
+
+# Bundle filenames. TABLE/HEAD are the mini ONNX graphs and must match what
+# tools/release.py publishes (research-side).
+TABLE = "mini-table-int8.onnx"
+HEAD = "mini-head-fABC.onnx"
+TOKENIZER = "tokenizer.json"
+MANIFEST = "manifest.json"
+
+# Each mode's bundle as {logical name -> filename}. Remote layout is <mode>/<filename>.
+BUNDLE_FILES = {
+    "mini": {"table": TABLE, "head": HEAD, "tokenizer": TOKENIZER, "manifest": MANIFEST},
+    "base": {"config": "config.json", "weights": "weights.pt",
+             "tokenizer": TOKENIZER, "manifest": MANIFEST},
+}
+
+
+@dataclass
+class Manifest:
+    """The per-model config that travels in the bundle (feats spec + z-score stats +
+    sizes), read from manifest.json so the package needs no torch to load it."""
+    mode: str
+    feats: str | None = None
+    hidden: int | None = None
+    window: int = 8192
+    cap: int = 0
+    zscore_mean: list = field(default_factory=list)
+    zscore_std: list = field(default_factory=list)
+
+    @classmethod
+    def from_file(cls, path):
+        d = json.loads(Path(path).read_text(encoding="utf-8"))
+        z = d.get("zscore") or {}
+        return cls(mode=d["mode"], feats=d.get("feats"), hidden=d.get("hidden"),
+                   window=d.get("window", 8192), cap=d.get("cap", 0),
+                   zscore_mean=z.get("mean", []), zscore_std=z.get("std", []))
+
+
+def resolve_bundle(mode, local_dir=None):
+    """{logical name -> local file path} for `mode`'s bundle. A local_dir reads the
+    files straight off disk (development); otherwise each is pulled from the release
+    repo at the pinned revision and cached by huggingface_hub."""
+    files = BUNDLE_FILES[mode]
+    if local_dir:
+        base = Path(local_dir)
+        return {name: base / fn for name, fn in files.items()}
+    return {name: Path(hf_hub_download(REPO, f"{mode}/{fn}", revision=REVISION))
+            for name, fn in files.items()}

htmlsift/cli.py

@@ -0,0 +1,29 @@
+"""`htmlsift` command line: prefetch model bundles into the Hugging Face cache.
+
+    htmlsift download mini        # fetch + cache the mini bundle
+    htmlsift download base        # fetch + cache the base bundle
+"""
+import argparse
+
+from ._artifacts import BUNDLE_FILES, resolve_bundle
+
+
+def _download(mode):
+    paths = resolve_bundle(mode)          # hf_hub_download each file -> cached
+    print(f"{mode}: cached {len(paths)} files")
+    for name, p in paths.items():
+        print(f"  {name:10} {p}")
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="htmlsift")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    dl = sub.add_parser("download", help="prefetch a model bundle into the HF cache")
+    dl.add_argument("mode", choices=list(BUNDLE_FILES))
+    args = ap.parse_args(argv)
+    if args.cmd == "download":
+        _download(args.mode)
+
+
+if __name__ == "__main__":
+    main()

htmlsift/core/__init__.py

@@ -0,0 +1,6 @@
+"""Vendored inference engine, relicensed Apache-2.0 from htmlsift-research/core.
+
+Kept as a faithful copy so each module diffs cleanly against its research original; the
+only intended change is render.py's output-only ``links=`` extension. See packaging_plan.md
+sections 4-5.
+"""

htmlsift/core/band.py

@@ -0,0 +1,144 @@
+"""
+Chunked band attention for the encoder's sliding-attention layers -- a speed
+optimization that changes nothing about the output.
+
+transformers builds the sliding window as a dense (B, 1, L, L) boolean mask and
+hands it to scaled_dot_product_attention, whose CPU kernel scores the whole
+L x L matrix and then discards the out-of-band entries -- so a sliding layer
+costs as much as a full one and local_attention buys nothing. This computes the
+same numbers in O(L*w) by cutting the sequence into chunks and letting each
+query chunk attend only to the key slice it can reach.
+
+The arithmetic is unchanged: the per-chunk predicate is exactly
+abs(q_idx - kv_idx) <= config.sliding_window, the same one transformers' sliding
+mask applies. So this is a speed-only optimization gated on numeric identity, not on
+a val score -- it computes the same numbers as stock sliding attention, only faster.
+
+    import core.band
+    core.band.enable(encoder)      # in place, idempotent
+"""
+import torch
+import torch.nn.functional as F
+from transformers.integrations.sdpa_attention import sdpa_attention_forward
+from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS, sdpa_mask
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+IMPL = "sdpa_band"
+CHUNK = 128
+# Below this the reshaping costs more than the scores it saves, and the dense
+# band mask is only a few hundred KB -- exactly what transformers would build.
+MIN_CHUNKED = 4 * CHUNK
+_NO_MASK = {"full_attention": None, "sliding_attention": None}
+
+
+def _dense_band_mask(n_tok, half, device):
+    """(1, 1, L, L) bool -- the sliding mask transformers would have built."""
+    i = torch.arange(n_tok, device=device)
+    return ((i[None, :] - i[:, None]).abs() <= half)[None, None]
+
+
+def _chunk_mask(chunk, half, n_chunks, n_tok, device):
+    """(n_chunks, 1, chunk, chunk + 2*half) bool: inside the band AND a real
+    token. A query row past the sequence end has no valid key and softmax over an
+    all-False row is NaN, so each dead row gets its own zero slot; those rows are
+    sliced off before the result returns."""
+    span = chunk + 2 * half
+    j = torch.arange(chunk, device=device)
+    t = torch.arange(span, device=device)
+    band = (t[None, :] - half - j[:, None]).abs() <= half
+    c = torch.arange(n_chunks, device=device)
+    kpos = c[:, None] * chunk - half + t[None, :]
+    m = band[None] & ((kpos >= 0) & (kpos < n_tok))[:, None, :]
+    dead = (c[:, None] * chunk + j[None, :]) >= n_tok
+    if bool(dead.any()):
+        ci, ji = torch.nonzero(dead, as_tuple=True)
+        m[ci, ji, half + ji] = True
+    return m[:, None]
+
+
+def band_sdpa_forward(module, query, key, value, attention_mask, dropout=0.0,
+                      scaling=None, sliding_window=None, **kwargs):
+    """Drop-in for sdpa_attention_forward. The band path runs only for a sliding
+    layer with nothing to pad; a full layer (sliding_window is None) or any real
+    padding mask is delegated to stock sdpa, unchanged."""
+    if sliding_window is None or attention_mask is not None:
+        return sdpa_attention_forward(module, query, key, value, attention_mask,
+                                      dropout=dropout, scaling=scaling,
+                                      sliding_window=sliding_window, **kwargs)
+
+    half = int(module.config.sliding_window)
+    B, H, L, D = query.shape
+
+    if L <= MIN_CHUNKED:
+        out = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=_dense_band_mask(L, half, query.device),
+            dropout_p=dropout, scale=scaling, is_causal=False)
+        return out.transpose(1, 2).contiguous(), None
+
+    chunk = max(CHUNK, half)
+    span = chunk + 2 * half
+    pad = (-L) % chunk
+    n = (L + pad) // chunk
+
+    q = F.pad(query, (0, 0, 0, pad)).view(B, H, n, chunk, D)
+    q = q.permute(0, 2, 1, 3, 4).reshape(B * n, H, chunk, D)
+    kv = []
+    for x in (key, value):
+        x = F.pad(x, (0, 0, half, pad + half)).unfold(2, span, chunk)
+        kv.append(x.permute(0, 2, 1, 4, 3).reshape(B * n, H, span, D))
+
+    m = _chunk_mask(chunk, half, n, L, query.device)
+    if B > 1:
+        m = m.repeat(B, 1, 1, 1)
+
+    out = F.scaled_dot_product_attention(q, kv[0], kv[1], attn_mask=m,
+                                         dropout_p=dropout, scale=scaling,
+                                         is_causal=False)
+    out = out.view(B, n, H, chunk, D).permute(0, 2, 1, 3, 4)
+    out = out.reshape(B, H, n * chunk, D)[:, :, :L]
+    return out.transpose(1, 2).contiguous(), None
+
+
+def register():
+    ALL_ATTENTION_FUNCTIONS.register(IMPL, band_sdpa_forward)
+    # Only reached on the fallback path (a real padding mask), but a missing key
+    # here would be a KeyError deep inside mask construction.
+    ALL_MASK_ATTENTION_FUNCTIONS.register(IMPL, sdpa_mask)
+
+
+def _no_mask_hook(module, args, kwargs):
+    """Replace an all-ones 2D mask with the "no mask" mapping so the model never
+    materialises the dense (1, 1, L, L) sliding mask -- the very thing the band
+    path exists to avoid. ModernBertModel.forward uses a mask dict verbatim when
+    given one, so {'full_attention': None, 'sliding_attention': None} makes every
+    layer receive None and band_sdpa_forward take the fast path. Anything that is
+    NOT an all-ones 2D mask (real padding) is left alone: the dense mask is built
+    and band_sdpa_forward falls back to stock sdpa. The fast path can therefore
+    only run when there is nothing to pad.
+
+    A forward_pre_hook, NOT a wrapped encoder.forward, and that is load-bearing.
+    torch.ao.quantization.quantize_dynamic returns a deep copy; a closure stored
+    as encoder.forward keeps pointing at the original fp32 module, so enable()
+    then quantize would silently run fp32 while reporting every layer quantized.
+    Hooks live in _forward_pre_hooks, are copied with the module, and take module
+    as an argument instead of capturing it.
+    """
+    am = kwargs.get("attention_mask")
+    if am is None or (torch.is_tensor(am) and am.dim() == 2 and bool(am.all())):
+        kwargs["attention_mask"] = _NO_MASK
+        return args, kwargs
+    return None
+
+
+def enable(encoder):
+    """Switch a ModernBERT encoder to band attention, in place. Idempotent."""
+    register()
+    encoder.config._attn_implementation = IMPL
+    for m in encoder.modules():          # layers may hold their own config copy
+        cfg = getattr(m, "config", None)
+        if cfg is not None:
+            cfg._attn_implementation = IMPL
+    if not getattr(encoder, "_band_enabled", False):
+        encoder.register_forward_pre_hook(_no_mask_hook, with_kwargs=True)
+        encoder._band_enabled = True
+    return encoder

htmlsift/core/features.py

@@ -0,0 +1,129 @@
+"""Structural block features for the table-embedding arm (mini): the signal a
+mean-pooled bag of token embeddings cannot carry, and nothing the rendered text
+already shows.
+
+Computed from render_tree's line_src (the DOM elements behind each block) -- no
+second parse, no marker injection, no change to the renderer. Booleans for
+presence; only the two genuine scalars (depth, log_chars) are z-scored, using the
+stats shipped in the model bundle.
+
+A column earns its place under one rule: it is absent from the tokenized block text
+AND not recoverable by the BiGRU. Tag structure is flattened by markdown, and any
+marker that survives is diluted by the per-block mean-pool; DOM depth and
+link-ancestry never appear in the rendered text at all (the renderer prints an <a>
+as bare text); block length is destroyed the moment the tokens are averaged.
+
+This module is the feature manifest: column order and the group map live here and
+nowhere else.
+"""
+
+import math
+
+import numpy as np
+
+# Group A -- ancestor-or-self tag presence, 30 booleans. Semantic containers that
+# markdown flattens away, plus block tags whose surviving text marker the mean-pool
+# dilutes.
+TAG_VOCAB = [
+    "main", "article", "section", "nav", "header", "footer", "aside",
+    "h1", "h2", "h3", "h4", "h5", "h6", "p", "li", "ul", "ol", "dl",
+    "table", "tr", "td", "th", "blockquote", "pre", "code",
+    "figure", "figcaption", "form", "button", "div",
+]
+
+GROUPS = {
+    "A": [f"tag_{t}" for t in TAG_VOCAB],   # ancestor/self tag presence (boolean)
+    "B": ["depth", "has_link"],             # DOM nesting depth (scalar); an <a> ancestor (boolean)
+    "C": ["log_chars"],                     # block length -- what the mean-pool averages away
+}
+
+FEATURE_NAMES = [name for g in ("A", "B", "C") for name in GROUPS[g]]
+
+# The only unbounded columns; z-scored with the bundle's stats. Booleans and
+# already-bounded columns stay raw.
+ZSCORE_COLS = ["depth", "log_chars"]
+
+_ORDER = ("A", "B", "C")
+_VOCAB = set(TAG_VOCAB)
+
+
+def collect_features(lines, line_src, groups=_ORDER):
+    """(lines, line_src) from render_tree -> [n_blocks, K] float32, one row per
+    non-empty line (the block unit). Columns are the requested groups in manifest
+    order; only those groups are computed, so a text-only arm passes groups=() and
+    pays nothing. The ancestor chain of each element is cached and shared across
+    blocks (below)."""
+    groups = [g for g in _ORDER if g in groups]
+    cols = [c for g in groups for c in GROUPS[g]]
+    col = {c: i for i, c in enumerate(cols)}
+    blocks = [(ln, src) for ln, src in zip(lines, line_src, strict=True) if ln.strip()]
+    out = np.zeros((len(blocks), len(cols)), dtype=np.float32)
+    need_a, need_b, need_c = "A" in groups, "B" in groups, "C" in groups
+
+    # element -> (vocab tags, depth, link) for its ancestor-or-self chain, computed
+    # once and reused. Every block re-reaches the same container chain to root, so
+    # caching collapses that repetition: the walk stops at the first cached ancestor.
+    # Keyed by the element object -- the cache holds the reference, which keeps the
+    # lxml proxy (and thus its identity) stable for the call.
+    chain = {}
+
+    def chain_info(e):
+        stack = []
+        node = e
+        while node is not None and isinstance(node.tag, str) and node not in chain:
+            stack.append(node)
+            node = node.getparent()
+        if node is not None and isinstance(node.tag, str):
+            tags, depth, link = chain[node]
+        else:
+            tags, depth, link = frozenset(), 0, False
+        for nd in reversed(stack):                     # fold from cached base to e
+            t = nd.tag
+            if need_a and t in _VOCAB and t not in tags:
+                tags = tags | {t}
+            depth += 1
+            link = link or t == "a"
+            chain[nd] = (tags, depth, link)
+        return tags, depth, link
+
+    for bi, (text, src) in enumerate(blocks):
+        if need_a or need_b:
+            btags, depth_sum, n_src, has_link = set(), 0, 0, False
+            for e in src:
+                et, ed, el = chain_info(e)
+                btags |= et
+                depth_sum += ed
+                n_src += 1
+                has_link = has_link or el
+            if need_a:
+                for t in btags:
+                    out[bi, col[f"tag_{t}"]] = 1.0
+            if need_b:
+                out[bi, col["depth"]] = depth_sum / n_src if n_src else 0.0
+                out[bi, col["has_link"]] = float(has_link)
+        if need_c:
+            out[bi, col["log_chars"]] = math.log1p(len(text))
+    return out
+
+
+def feature_dim(groups):
+    """Number of feature columns a group selection produces (None/"" -> 0). Sizes
+    the head's d_in without hardcoding K."""
+    return sum(len(GROUPS[g]) for g in _ORDER if groups and g in groups)
+
+
+def _zscore_idx(groups):
+    """Column positions of the unbounded scalars (ZSCORE_COLS) within a group
+    selection's layout -- lets the loader pair the bundle's mean/std with the right
+    columns for a given feats group."""
+    cols = [c for g in _ORDER if g in groups for c in GROUPS[g]]
+    return [i for i, c in enumerate(cols) if c in ZSCORE_COLS]
+
+
+def apply_zscore(fa, stats):
+    """A copy of a page's [n_blocks, K] array with the z-scored columns normalized by
+    the bundle's (idx, mean, std) stats; booleans and bounded columns untouched."""
+    idx, mean, std = stats
+    out = fa.copy()
+    out[:, idx] = (out[:, idx] - mean) / std
+    return out

htmlsift/core/onnx_infer.py

@@ -0,0 +1,38 @@
+"""Torch-free onnxruntime forward for the mini model (int8 table + BiGRU head).
+
+Embed the token ids through the int8 table graph, mean-pool by block membership,
+concat any structural features, run the head graph, sigmoid. numpy + onnxruntime only
+-- no torch. Lifted unchanged from the research export so it stays bit-identical to the
+graphs' parity-gated reference forward.
+"""
+import numpy as np
+import onnxruntime as ort
+
+
+def pool(emb, members):
+    """Token embeddings emb [T, H] -> per-block mean vectors [n, H]. A token-less block
+    pools to zeros (sum of nothing / 1)."""
+    pooled = np.zeros((len(members), emb.shape[1]), dtype=np.float32)
+    for i, mem in enumerate(members):
+        if mem:
+            pooled[i] = emb[mem].mean(0)
+    return pooled
+
+
+def load_mini(table_path, head_path, threads=None):
+    """The two mini graphs -> infer_fn(page) -> per-block probs (np.float32): the
+    (page -> probs) contract the extractor consumes."""
+    opts = ort.SessionOptions()
+    if threads:
+        opts.intra_op_num_threads, opts.inter_op_num_threads = threads, 1
+    table = ort.InferenceSession(str(table_path), opts, providers=["CPUExecutionProvider"])
+    head = ort.InferenceSession(str(head_path), opts, providers=["CPUExecutionProvider"])
+
+    def infer_fn(page):
+        emb = table.run(["out"], {"ids": page["ids"].astype(np.int64)})[0]     # [T, H]
+        x = pool(emb, page["members"])                                         # [n, H]
+        if page.get("feats") is not None:
+            x = np.concatenate([x, page["feats"].astype(np.float32)], 1)       # [n, H+K]
+        logits = head.run(["logits"], {"x": x[None]})[0][0]                    # [n]
+        return (1.0 / (1.0 + np.exp(-logits))).astype(np.float32)
+    return infer_fn

htmlsift/core/prep_page.py

@@ -0,0 +1,93 @@
+"""Page layout for inference: tokenize a page's blocks once, map every token back
+to the block it came from, and (for the windowed encoder path) plan the token
+windows. No gradients, no encoder -- deterministic shaping the forwards consume.
+"""
+
+import numpy as np
+
+
+def block_spans(blocks):
+    """Char span [start, end) of each block within "\\n".join(blocks)."""
+    spans, pos = [], 0
+    for b in blocks:
+        spans.append((pos, pos + len(b)))
+        pos += len(b) + 1                     # +1 for the joining newline
+    return spans
+
+
+def page_members(enc, blocks):
+    """Per block, the indices of the tokens whose characters fall inside it.
+
+    `enc` is a fast-tokenizer output carrying offset_mapping and
+    special_tokens_mask. A single cursor advances through the non-special tokens as
+    blocks are consumed left to right, so the whole map is one linear pass. A token
+    joins a block when it overlaps the block's characters; the `end > bs` guard
+    drops zero-width tokens (e.g. (0,0)) while keeping a subword that straddles the
+    joining newline with the block it reaches into. A block with no tokens of its
+    own gets [] and pools to a zero vector, so members stays aligned with blocks.
+    """
+    offsets = enc["offset_mapping"]
+    special = enc["special_tokens_mask"]
+    tok_idx = [i for i in range(len(offsets)) if not special[i]]
+    members = []
+    ti = 0
+    for bs, be in block_spans(blocks):
+        while ti < len(tok_idx) and offsets[tok_idx[ti]][1] <= bs:
+            ti += 1
+        j, mem = ti, []
+        while j < len(tok_idx) and offsets[tok_idx[j]][0] < be:
+            if offsets[tok_idx[j]][1] > bs:
+                mem.append(tok_idx[j])
+            j += 1
+        members.append(mem)
+    return members
+
+
+def window_starts(n, window, stride=None):
+    """Start offsets of the windows tiling n tokens (n > window). stride defaults to
+    window // 2 (50% overlap); the last window snaps back to n - window so the tail
+    is always fully covered. Shared with the encoder stitch (torch_infer.pool_page),
+    so window boundaries are defined in one place."""
+    stride = stride or window // 2
+    starts = list(range(0, n - window + 1, stride))
+    if not starts:
+        starts = [0]
+    if starts[-1] != n - window:
+        starts.append(n - window)
+    return starts
+
+
+def apply_cap(ids, members, k):
+    """Keep each block's first k member tokens, dropping the rest from the sequence
+    -- a per-block width cap. Tokens belonging to no block (specials, the newline
+    joins) are always kept. Members are renumbered onto the compacted ids."""
+    keep = np.ones(len(ids), dtype=bool)
+    for mem in members:
+        if len(mem) > k:
+            keep[np.asarray(mem[k:], dtype=np.int64)] = False
+    newpos = np.cumsum(keep) - 1
+    return ids[keep], [[int(newpos[t]) for t in mem[:k]] for mem in members]
+
+
+def prep_page(tok, blocks, window, stride=None, cap=0, feats=None, tid=None):
+    """Blocks -> the page dict the forwards consume: token ids, per-block members,
+    optional structural features, and the window plan. The page carries its own
+    window/stride so the encoder stitch cannot drift. truncation=False -- the page is
+    windowed by the encoder, never truncated by the tokenizer. cap>0 applies the
+    per-block width cap first.
+
+    feats, when given, is an [n_blocks, K] array (already z-scored on the training
+    stats, aligned with blocks) carried through to the head concat; feats=None runs
+    text-only. The width cap drops tokens, not blocks, so feats stays aligned."""
+    enc = tok("\n".join(blocks), add_special_tokens=True,
+              return_offsets_mapping=True, return_special_tokens_mask=True,
+              return_tensors=None, truncation=False)
+    ids = np.array(enc["input_ids"], dtype=np.int64)
+    members = page_members(enc, blocks)
+    if cap:
+        ids, members = apply_cap(ids, members, cap)
+    if feats is not None:
+        feats = np.asarray(feats, dtype=np.float32)
+        assert feats.shape[0] == len(members), (feats.shape[0], len(members))
+    return {"tid": tid, "ids": ids, "members": members, "feats": feats,
+            "window": window, "stride": stride}

htmlsift/core/prune.py

@@ -0,0 +1,42 @@
+"""Prune a parsed DOM to the elements behind the selected blocks, then serialize.
+
+The renderer records, per rendered line, the DOM elements that produced it (line_src).
+Given the elements behind the KEPT lines, keep exactly those, their ancestors (the
+structural path to root) and their descendants (inline links/images inside a kept block),
+and drop everything else. lxml only -- no html2text.
+"""
+from lxml import html as lxml_html
+
+
+def keep_set(elements):
+    """The nodes to keep: each element, its ancestors, and its descendants. Holding the
+    proxies in the set keeps their lxml identity stable for the prune walk."""
+    keep = set()
+    for el in elements:
+        if el in keep:
+            continue
+        keep.add(el)
+        keep.update(el.iterancestors())
+        keep.update(el.iterdescendants())
+    return keep
+
+
+def prune(root, elements):
+    """Remove from root every element not in keep_set(elements); returns root (mutated)."""
+    keep = keep_set(elements)
+
+    def walk(el):
+        for child in list(el):
+            if child in keep:
+                walk(child)
+            else:
+                el.remove(child)
+
+    walk(root)
+    return root
+
+
+def to_html(root, elements):
+    """Prune root to the selected elements and serialize to an HTML string."""
+    prune(root, elements)
+    return lxml_html.tostring(root, encoding="unicode")