#!/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"(? 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