src/chunkers/base.py

#!/usr/bin/env python3
"""
Chunker framework: word-space boundaries, gold-block tracking, post-processing.

THREE DESIGN DECISIONS, EACH EVIDENCE-BACKED
--------------------------------------------
1. BOUNDARIES LIVE IN WORD SPACE, NOT TOKEN SPACE.
   Indic tokenizer fertility varies ~1.5-3x (measured by src/fertility.py). A
   fixed token budget yields different chunk counts per language, which silently
   breaks canonical cross-lingual chunk IDs. Word/sentence boundaries are
   language-invariant by construction. Token budgets are DERIVED per language
   from measured fertility, not assumed.

2. EVERY CHUNK RECORDS WHICH GOLD BLOCKS IT SPANS.
   On pseudo-documents the gold blocks are the original MS MARCO passages, so
   Block Integrity (arXiv:2603.25333) becomes exactly "did we split a passage?"
   -- a free, exact structural metric on a corpus with no native structure.

3. POST-PROCESSING IS MANDATORY, NOT OPTIONAL.
   Tiny-chunk merging + oversized re-splitting is worth +6 to +16 percentage
   points on mean intrinsic score (arXiv:2603.25333, sec 3.2) for ~30 lines of
   code. It is the cheapest win in the entire literature. It runs on EVERY
   strategy, including the baselines, so comparisons stay fair.

OVERLAP POLICY IS PER FAMILY, NOT GLOBAL.
   Overlap only compensates for arbitrary boundaries. High for fixed-size; low
   for semantic/structural (the boundary already follows meaning); NOT APPLICABLE
   for late chunking, where document context is already inside every chunk vector.
   Each Chunker declares `default_overlap` so this is enforced, not remembered.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Callable, Iterable, Protocol, Sequence

__all__ = [
    "Chunk", "Document", "Chunker", "REGISTRY", "register",
    "postprocess", "split_sentences", "words_of",
]

# ---------------------------------------------------------------- text utils

# Sentence terminators across the scripts in MSMARCO-XI:
#   .!?          Latin / most Indic (Indic scripts use Latin punctuation in practice)
#   ।       DEVANAGARI DANDA        -- Hindi, Marathi, Nepali, Sanskrit
#   ॥       DEVANAGARI DOUBLE DANDA
#   ۔       ARABIC FULL STOP        -- Urdu
#   ܁܂ SYRIAC (defensive)
_TERMINATORS = "।॥۔܁܂.!?"
_WS_RE = re.compile(r"\s+")

# Digits in every script present in MSMARCO-XI.
_DIGIT = "0-9०-९০-৯੦-੯૦-૯୦-୯௦-௯౦-౯೦-೯൦-൯٠-٩۰-۹"

# A terminator ends a sentence UNLESS it is a decimal point sitting between
# digits.  Without the guard, "Average low of 73.3 degrees" splits into three
# "sentences" -- which silently shreds every NUMERIC answer in the corpus, and
# NUMERIC is 25% of queries.  Found by the reader evaluation, not by inspection.
_SPLIT_RE = re.compile(
    rf"(?<![{_DIGIT}])[{_TERMINATORS}]+\s*"      # terminator not preceded by a digit
    rf"|[{_TERMINATORS}]+(?![{_DIGIT}])\s*"      # or not followed by one
)


def split_sentences(text: str) -> list[str]:
    """Script-aware sentence split that does not break decimals.

    Handles Latin (.!?), Devanagari danda (।॥) and Urdu full stop (۔).
    """
    if not text.strip():
        return []
    out, last = [], 0
    for m in _SPLIT_RE.finditer(text):
        piece = text[last:m.end()].strip()
        if piece:
            out.append(piece)
        last = m.end()
    tail = text[last:].strip()
    if tail:
        out.append(tail)
    return out or [text.strip()]


def words_of(text: str) -> list[str]:
    """Whitespace words. Language-invariant unit for boundary arithmetic."""
    return _WS_RE.split(text.strip()) if text.strip() else []


# ---------------------------------------------------------------- data model

@dataclass(slots=True)
class Document:
    """A pseudo-document (or a single passage, when corpus_form='raw')."""
    doc_id: str
    lang: str
    text: str
    # Gold blocks = original passages. (start_word, end_word, canonical_id) half-open.
    blocks: list[tuple[int, int, str]] = field(default_factory=list)

    @property
    def words(self) -> list[str]:
        return words_of(self.text)

    @classmethod
    def from_blocks(cls, doc_id: str, lang: str, texts: Sequence[str],
                    block_ids: Sequence[str], sep: str = "\n\n") -> "Document":
        """Build a document from its passages, recording exact word offsets."""
        blocks, cursor, parts = [], 0, []
        for t, bid in zip(texts, block_ids):
            n = len(words_of(t))
            if n == 0:
                continue
            blocks.append((cursor, cursor + n, bid))
            cursor += n
            parts.append(t)
        return cls(doc_id=doc_id, lang=lang, text=sep.join(parts), blocks=blocks)


@dataclass(slots=True)
class Chunk:
    doc_id: str
    lang: str
    text: str
    start_word: int
    end_word: int          # half-open
    strategy: str
    block_ids: list[str] = field(default_factory=list)   # gold blocks overlapped
    split_blocks: int = 0                                # gold blocks cut mid-way
    meta: dict = field(default_factory=dict)

    @property
    def n_words(self) -> int:
        return self.end_word - self.start_word

    def chunk_id(self) -> str:
        """Deterministic and language-independent in word space, so the SAME chunk
        in Hindi and Tamil carries the SAME id -- the basis for cross-lingual fusion."""
        return f"{self.doc_id}:{self.start_word}-{self.end_word}:{self.strategy}"


def _attach_blocks(chunk: Chunk, blocks: list[tuple[int, int, str]]) -> Chunk:
    """Record overlapped gold blocks and count how many were cut."""
    ids, split = [], 0
    for b0, b1, bid in blocks:
        if b1 <= chunk.start_word or b0 >= chunk.end_word:
            continue
        ids.append(bid)
        if b0 < chunk.start_word or b1 > chunk.end_word:
            split += 1
    chunk.block_ids, chunk.split_blocks = ids, split
    return chunk


def make_chunk(doc: Document, words: list[str], start: int, end: int,
               strategy: str, **meta) -> Chunk:
    start, end = max(0, start), min(len(words), end)
    c = Chunk(
        doc_id=doc.doc_id, lang=doc.lang, text=" ".join(words[start:end]),
        start_word=start, end_word=end, strategy=strategy, meta=meta,
    )
    return _attach_blocks(c, doc.blocks)


# ---------------------------------------------------------------- protocol

class Chunker(Protocol):
    name: str
    family: str
    default_overlap: float      # fraction of chunk size; 0.0 where overlap is meaningless

    def chunk(self, doc: Document) -> list[Chunk]: ...


REGISTRY: dict[str, type] = {}


def register(cls):
    REGISTRY[cls.name] = cls
    return cls


# ---------------------------------------------------------------- post-processing

def postprocess(
    chunks: list[Chunk],
    doc: Document,
    min_words: int,
    max_words: int,
    merge_ceiling: float = 1.05,
) -> list[Chunk]:
    """Tiny-chunk merge + oversized re-split.  arXiv:2603.25333 sec 3.2: +6..+16 pp.

    min_words / max_words are WORD counts derived per language from measured
    fertility -- never hard-coded token counts.

    Raw chunkers emit near-empty fragments (the paper reports 0-4 token chunks)
    that waste top-k slots and dilute embeddings. This pass removes them without
    disturbing well-formed segments.
    """
    if not chunks:
        return chunks
    words = doc.words
    ceiling = int(max_words * merge_ceiling)

    # -- pass 1: merge tiny chunks into the smaller adjacent neighbour ---------
    merged: list[Chunk] = []
    for c in sorted(chunks, key=lambda x: x.start_word):
        if merged and c.n_words < min_words:
            prev = merged[-1]
            if (c.end_word - prev.start_word) <= ceiling:
                merged[-1] = make_chunk(doc, words, prev.start_word, c.end_word,
                                        c.strategy, **{**prev.meta, "merged": True})
                continue
        merged.append(c)

    # a tiny FIRST chunk has no left neighbour: fold it right
    if len(merged) > 1 and merged[0].n_words < min_words:
        a, b = merged[0], merged[1]
        if (b.end_word - a.start_word) <= ceiling:
            merged[1] = make_chunk(doc, words, a.start_word, b.end_word,
                                   b.strategy, **{**b.meta, "merged": True})
            merged.pop(0)

    # -- pass 2: re-split oversized chunks on sentence boundaries --------------
    out: list[Chunk] = []
    for c in merged:
        if c.n_words <= max_words:
            out.append(c)
            continue
        cuts = _sentence_cuts(c.text, c.start_word, max_words)
        prev = c.start_word
        for cut in cuts + [c.end_word]:
            if cut > prev:
                out.append(make_chunk(doc, words, prev, cut, c.strategy,
                                      **{**c.meta, "resplit": True}))
                prev = cut
    return out


def _sentence_cuts(text: str, offset: int, max_words: int) -> list[int]:
    """Word indices at which to cut an oversized chunk, preferring sentence ends."""
    cuts, cursor, since = [], offset, 0
    for sent in split_sentences(text):
        n = len(words_of(sent))
        if since + n > max_words and since > 0:
            cuts.append(cursor)
            since = 0
        cursor += n
        since += n
    return cuts


# ---------------------------------------------------------------- budgets

def word_budget(target_tokens: int, fertility: float, en_tokens_per_word: float = 1.35
                ) -> tuple[int, int]:
    """Convert a token intent into a language-specific WORD budget.

    Example: target 200 tokens, Hindi fertility 2.1 vs English
        english words  = 200 / 1.35            ~= 148
        hindi   words  = 200 / (1.35 * 2.1)    ~=  70

    Both produce ~200 tokens for the embedder, but the WORD spans differ -- which
    is exactly right, because the same content occupies fewer words in a script
    the tokenizer fragments more. Chunk boundaries stay aligned to content, so
    canonical cross-lingual chunk IDs survive.
    """
    tpw = max(0.2, en_tokens_per_word * max(0.1, fertility))
    target = max(8, int(round(target_tokens / tpw)))
    return max(4, int(target * 0.25)), int(target * 1.6)   # (min_words, max_words)


def chunk_document(doc: Document, chunker: Chunker, min_words: int, max_words: int,
                   post: bool = True) -> list[Chunk]:
    chunks = chunker.chunk(doc)
    return postprocess(chunks, doc, min_words, max_words) if post else chunks


def chunk_corpus(docs: Iterable[Document], chunker: Chunker,
                 budget_fn: Callable[[str], tuple[int, int]],
                 post: bool = True) -> list[Chunk]:
    """budget_fn maps language code -> (min_words, max_words) from fertility.json."""
    out: list[Chunk] = []
    for d in docs:
        lo, hi = budget_fn(d.lang)
        out.extend(chunk_document(d, chunker, lo, hi, post))
    return out