chunk.py

"""Inference-side semantic chunker.

The training task is "does a segment boundary follow this sentence"; this module
turns that into the thing a RAG pipeline actually wants: a list of text chunks,
subject to hard size limits an embedder can accept.

Two things here are not in the training loop and matter in production:

1. Windowed inference. Real documents run far past 2048 tokens. We slide a
   window with sentence-level left context so boundary decisions near a window
   edge still see what came before, instead of being scored cold.
2. Hard size caps. A semantic boundary model will happily emit a 4000-token
   chunk if the topic does not shift. Embedders truncate at their own limit, so
   `max_chunk_tokens` force-splits at the lowest-confidence interior sentence
   rather than at an arbitrary character offset.
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from typing import Callable, Iterable

import numpy as np
from transformers import AutoTokenizer

#: Published weights. Used when no model path is given, so `SemanticChunker()`
#: works with nothing downloaded by hand.
DEFAULT_MODEL = "0xKitkat/semantic-chunker-modernbert-base"

#: Where the ONNX graph lives inside that repo.
ONNX_SUBFOLDER = "onnx"

# Sentence splitter: deliberately dependency-free. Handles the common
# abbreviation cases that naive `.` splitting gets wrong. Pass your own via
# `sentence_splitter=` if you already have spacy/nltk/pysbd in the pipeline.
_ABBREV = (
    r"(?<!\bMr)(?<!\bMrs)(?<!\bMs)(?<!\bDr)(?<!\bProf)(?<!\bSt)(?<!\bJr)(?<!\bSr)"
    r"(?<!\bInc)(?<!\bLtd)(?<!\bCo)(?<!\bvs)(?<!\betc)(?<!\bi\.e)(?<!\be\.g)"
    r"(?<!\bFig)(?<!\bNo)(?<!\bVol)(?<!\bApprox)(?<!\b[A-Z])"
)
_SENT_RE = re.compile(rf"{_ABBREV}(?<=[.!?])[\"')\]]*\s+(?=[A-Z0-9\"'(\[])")


def default_sentence_splitter(text: str) -> list[str]:
    out = []
    for block in re.split(r"\n\s*\n", text):  # paragraph breaks are always splits
        block = block.strip()
        if not block:
            continue
        out.extend(s.strip() for s in _SENT_RE.split(block) if s.strip())
    return out


def _softmax_last(x: np.ndarray) -> np.ndarray:
    x = x - x.max(axis=-1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)


class TorchBackend:
    """Default backend. torch is imported here so the ONNX path stays torch-free."""

    def __init__(self, model_path: str, device: str | None = None):
        import torch
        from transformers import AutoModelForTokenClassification

        self._torch = torch
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = (
            AutoModelForTokenClassification.from_pretrained(model_path)
            .eval()
            .to(self.device)
        )
        self.id2label = dict(self.model.config.id2label)

    def __call__(self, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
        t = self._torch
        with t.no_grad():
            out = self.model(
                input_ids=t.tensor(input_ids, dtype=t.long, device=self.device),
                attention_mask=t.tensor(attention_mask, dtype=t.long, device=self.device),
            ).logits
        return out.float().cpu().numpy()


def resolve_onnx_dir(path_or_repo: str, subfolder: str = ONNX_SUBFOLDER) -> str:
    """Accept a local ONNX dir, a local repo root, or a Hub repo id.

    Without this, ONNX users have to know that the graph sits in a subfolder and
    fetch it themselves before anything works -- which is exactly the friction
    the ONNX build exists to remove.
    """
    if os.path.isdir(path_or_repo):
        for cand in (path_or_repo, os.path.join(path_or_repo, subfolder)):
            if os.path.exists(os.path.join(cand, "model.onnx")):
                return cand
        raise FileNotFoundError(
            f"no model.onnx in {path_or_repo} or its {subfolder}/ subfolder"
        )

    from huggingface_hub import snapshot_download

    local = snapshot_download(path_or_repo, allow_patterns=f"{subfolder}/*")
    return os.path.join(local, subfolder)


class OnnxBackend:
    """CPU inference with onnxruntime only -- no torch in the dependency tree."""

    def __init__(self, model_dir: str, providers: list[str] | None = None):
        import json
        import onnxruntime as ort

        model_dir = resolve_onnx_dir(model_dir)
        path = os.path.join(model_dir, "model.onnx")
        if not os.path.exists(path):
            raise FileNotFoundError(f"no model.onnx under {model_dir}")
        self.sess = ort.InferenceSession(
            path, providers=providers or ["CPUExecutionProvider"]
        )
        with open(os.path.join(model_dir, "config.json"), encoding="utf-8") as f:
            cfg = json.load(f)
        self.id2label = cfg.get("id2label", {0: "O", 1: "semantic-shift"})

    def __call__(self, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
        return self.sess.run(
            None,
            {
                "input_ids": input_ids.astype(np.int64),
                "attention_mask": attention_mask.astype(np.int64),
            },
        )[0]


@dataclass
class Chunk:
    text: str
    start_sentence: int
    end_sentence: int
    n_tokens: int
    boundary_score: float  # confidence of the boundary that closed this chunk


class SemanticChunker:
    def __init__(
        self,
        model_path: str | None = None,
        threshold: float = 0.5,
        max_chunk_tokens: int = 512,
        min_chunk_tokens: int = 0,
        min_chunk_sentences: int = 1,
        window_tokens: int = 2048,
        context_sentences: int = 3,
        device: str | None = None,
        sentence_splitter: Callable[[str], list[str]] | None = None,
        batch_size: int = 8,
        backend: str = "torch",
        tokenizer_path: str | None = None,
    ):
        model_path = model_path or DEFAULT_MODEL
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path or model_path)
        if backend == "onnx":
            self.backend = OnnxBackend(model_path)
        elif backend == "torch":
            self.backend = TorchBackend(model_path, device)
        else:
            raise ValueError(f"unknown backend {backend!r}, expected 'torch' or 'onnx'")
        self.device = getattr(self.backend, "device", "cpu")
        self._configure(
            threshold=threshold,
            max_chunk_tokens=max_chunk_tokens,
            min_chunk_tokens=min_chunk_tokens,
            min_chunk_sentences=min_chunk_sentences,
            window_tokens=window_tokens,
            context_sentences=context_sentences,
            sentence_splitter=sentence_splitter,
            batch_size=batch_size,
        )

    @classmethod
    def from_backend(cls, backend, tokenizer, **kwargs):
        """Wrap a model somebody else already loaded.

        `transformers.pipeline` owns model construction -- it resolves the
        checkpoint, picks the device, and applies dtype/quantisation settings.
        Re-loading the weights underneath it would double the memory and
        silently ignore all of that, so the pipeline builds the chunker around
        the model it was handed instead.
        """
        self = cls.__new__(cls)
        self.tokenizer = tokenizer
        self.backend = backend
        self.device = getattr(backend, "device", "cpu")
        self._configure(**kwargs)
        return self

    def _configure(
        self,
        threshold: float = 0.5,
        max_chunk_tokens: int = 512,
        min_chunk_tokens: int = 0,
        min_chunk_sentences: int = 1,
        window_tokens: int = 2048,
        context_sentences: int = 3,
        sentence_splitter: Callable[[str], list[str]] | None = None,
        batch_size: int = 8,
    ):
        self.threshold = threshold
        self.max_chunk_tokens = max_chunk_tokens
        self.min_chunk_tokens = min_chunk_tokens
        self.min_chunk_sentences = max(1, min_chunk_sentences)
        self.window_tokens = window_tokens
        self.context_sentences = context_sentences
        self.split_sentences = sentence_splitter or default_sentence_splitter
        self.batch_size = batch_size

        id2label = dict(self.backend.id2label)
        self.pos_idx = next(
            (int(i) for i, n in id2label.items()
             if str(n).lower() in {"semantic-shift", "separator", "shift"}),
            1,
        )

    # ------------------------------------------------------------------
    def boundary_scores(self, sentences: list[str]) -> list[float]:
        """P(boundary follows) for each sentence. Last entry is always 1.0."""
        if not sentences:
            return []
        if len(sentences) == 1:
            return [1.0]

        toks = self.tokenizer(
            [s if i == 0 else " " + s for i, s in enumerate(sentences)],
            add_special_tokens=False,
        )["input_ids"]

        budget = self.window_tokens - 2
        windows = []  # (ids, [(sentence_idx, token_pos)])
        start = 0
        n = len(sentences)
        while start < n:
            ids, marks = [], []
            ctx_start = max(0, start - self.context_sentences)
            for i in range(ctx_start, start):
                ids.extend(toks[i])
            i = start
            while i < n:
                t = toks[i] or [self.tokenizer.unk_token_id]
                if len(t) > budget:
                    t = t[:budget]
                if len(ids) + len(t) > budget:
                    break
                ids.extend(t)
                marks.append((i, len(ids) - 1))
                i += 1
            if i == start:  # pathological single sentence
                marks.append((start, max(len(ids) - 1, 0)))
                i = start + 1
            windows.append((ids, marks))
            start = i

        scores = [0.0] * n
        for b in range(0, len(windows), self.batch_size):
            batch = windows[b : b + self.batch_size]
            maxlen = max(len(w[0]) for w in batch) + 2
            input_ids = np.full(
                (len(batch), maxlen), self.tokenizer.pad_token_id, dtype=np.int64
            )
            attn = np.zeros((len(batch), maxlen), dtype=np.int64)
            for r, (ids, _) in enumerate(batch):
                seq = [self.tokenizer.cls_token_id] + ids + [self.tokenizer.sep_token_id]
                input_ids[r, : len(seq)] = seq
                attn[r, : len(seq)] = 1
            logits = self.backend(input_ids, attn)
            probs = _softmax_last(logits.astype(np.float32))[..., self.pos_idx]
            for r, (_, marks) in enumerate(batch):
                for sent_idx, tok_pos in marks:
                    scores[sent_idx] = float(probs[r, tok_pos + 1])  # +1 for CLS

        scores[-1] = 1.0  # document end is always a boundary
        return scores

    # ------------------------------------------------------------------
    def _token_lens(self, texts: list[str]) -> list[int]:
        """One batched tokenizer call, not one per sentence."""
        if not texts:
            return []
        enc = self.tokenizer(texts, add_special_tokens=False)["input_ids"]
        return [len(x) for x in enc]

    def split(self, text: str) -> list[Chunk]:
        sentences = self.split_sentences(text)
        if not sentences:
            return []
        scores = self.boundary_scores(sentences)
        sent_tokens = self._token_lens(sentences)

        chunks: list[Chunk] = []
        buf_start = 0
        buf_tokens = 0

        def emit(end_idx: int, score: float):
            nonlocal buf_start, buf_tokens
            body = " ".join(sentences[buf_start : end_idx + 1]).strip()
            if body:
                chunks.append(
                    Chunk(
                        text=body,
                        start_sentence=buf_start,
                        end_sentence=end_idx,
                        n_tokens=sum(sent_tokens[buf_start : end_idx + 1]),
                        boundary_score=score,
                    )
                )
            buf_start = end_idx + 1
            buf_tokens = 0

        for i, sent in enumerate(sentences):
            buf_tokens += sent_tokens[i]
            n_in_buf = i - buf_start + 1

            # hard cap wins over semantics: split at the weakest interior
            # boundary so we cut where the model is least confident, not
            # wherever the token counter happened to run out.
            if buf_tokens > self.max_chunk_tokens and n_in_buf > 1:
                interior = range(buf_start, i)
                cut = max(interior, key=lambda j: scores[j])
                emit(cut, scores[cut])
                buf_tokens = sum(sent_tokens[buf_start : i + 1])
                n_in_buf = i - buf_start + 1

            if scores[i] >= self.threshold and n_in_buf >= self.min_chunk_sentences:
                emit(i, scores[i])

        if buf_start < len(sentences):
            emit(len(sentences) - 1, 1.0)

        return self._merge_undersized(chunks, sentences, sent_tokens)

    def _merge_undersized(self, chunks, sentences, sent_tokens):
        """Merge chunks below min_chunk_tokens into their neighbour.

        Measured need, not speculation: on short-document corpora the boundary
        model happily splits a 150-token document in two, and two 75-token
        fragments retrieve worse than one whole document. A topic boundary is
        real there, but acting on it is counterproductive -- the embedder needs
        enough text to place the vector well.

        Merges forward by preference (keeps reading order), backward for a
        trailing runt, and never produces a chunk exceeding max_chunk_tokens.
        """
        if self.min_chunk_tokens <= 0 or len(chunks) < 2:
            return chunks

        out = []
        i = 0
        while i < len(chunks):
            cur = chunks[i]
            while (
                cur.n_tokens < self.min_chunk_tokens
                and i + 1 < len(chunks)
                and cur.n_tokens + chunks[i + 1].n_tokens <= self.max_chunk_tokens
            ):
                nxt = chunks[i + 1]
                cur = Chunk(
                    text=(cur.text + " " + nxt.text).strip(),
                    start_sentence=cur.start_sentence,
                    end_sentence=nxt.end_sentence,
                    n_tokens=cur.n_tokens + nxt.n_tokens,
                    boundary_score=nxt.boundary_score,
                )
                i += 1
            out.append(cur)
            i += 1

        # a trailing runt has no forward neighbour left; fold it back
        if (
            len(out) > 1
            and out[-1].n_tokens < self.min_chunk_tokens
            and out[-2].n_tokens + out[-1].n_tokens <= self.max_chunk_tokens
        ):
            last, prev = out.pop(), out.pop()
            out.append(
                Chunk(
                    text=(prev.text + " " + last.text).strip(),
                    start_sentence=prev.start_sentence,
                    end_sentence=last.end_sentence,
                    n_tokens=prev.n_tokens + last.n_tokens,
                    boundary_score=last.boundary_score,
                )
            )
        return out

    def split_text(self, text: str) -> list[str]:
        """Convenience: just the strings."""
        return [c.text for c in self.split(text)]

    def batch_split(self, texts: Iterable[str]) -> list[list[str]]:
        return [self.split_text(t) for t in texts]


def _cli(argv=None):
    import argparse
    import json as _json
    import sys

    ap = argparse.ArgumentParser(
        prog="boundary-chunk",
        description="Split a document into topic-coherent chunks.",
    )
    ap.add_argument("--model", default=None, help=f"default: {DEFAULT_MODEL}")
    ap.add_argument("--threshold", type=float, default=0.5)
    ap.add_argument("--max-chunk-tokens", type=int, default=512)
    ap.add_argument("--min-chunk-tokens", type=int, default=0)
    ap.add_argument("--file", help="read text from file; otherwise stdin")
    ap.add_argument("--backend", default="torch", choices=["torch", "onnx"])
    ap.add_argument("--tokenizer", default=None, help="only needed for a bare onnx dir")
    ap.add_argument("--json", action="store_true", help="emit JSON instead of text")
    args = ap.parse_args(argv)

    text = open(args.file, encoding="utf-8").read() if args.file else sys.stdin.read()
    ch = SemanticChunker(
        args.model,
        threshold=args.threshold,
        max_chunk_tokens=args.max_chunk_tokens,
        min_chunk_tokens=args.min_chunk_tokens,
        backend=args.backend,
        tokenizer_path=args.tokenizer,
    )
    chunks = ch.split(text)
    if args.json:
        print(_json.dumps([c.__dict__ for c in chunks], ensure_ascii=False, indent=2))
    else:
        for i, c in enumerate(chunks):
            print(f"\n--- chunk {i}  ({c.n_tokens} tok, score {c.boundary_score:.3f}) ---")
            print(c.text)
    return 0


if __name__ == "__main__":
    raise SystemExit(_cli())