chunking.py

"""Sentence-level chunking for late-interaction (MaxSim) scoring.

A whole-document embedding mean-pools a 25-sentence article into one vector, so
a topic named in a single sentence gets drowned out. Embedding each sentence
separately avoids that dilution: a query is scored against the best-matching
sentence chunk rather than the blurred average of the article.

This module is the pure, dependency-free segmenter that produces those chunks.
Chunk boundaries are part of the scoring contract — change them and previously
computed vectors no longer correspond to the same text.
"""

import re

# Below this length a sentence is treated as a fragment and glued onto the
# following one instead of becoming its own chunk.
MIN_CHUNK_CHARS = 25

# Upper bound on body chunks per article (title is always kept on top of
# this). Long articles are front-loaded — the lede carries the entities.
MAX_NEWS_CHUNKS = 60

# Abbreviations / patterns that would otherwise over-split news copy
# (U.S., Inc., $3.5B, Sept. 5, e.g.). Deliberately no NLP dependency — good
# enough for embedding-chunk boundaries, where a mis-split costs almost
# nothing.
ABBREVIATIONS = {
    "u.s", "u.k", "e.u", "inc", "corp", "co", "ltd", "llc", "plc",
    "mr", "mrs", "ms", "dr", "prof", "sen", "rep", "gov", "gen",
    "jan", "feb", "mar", "apr", "jun", "jul", "aug", "sept", "sep",
    "oct", "nov", "dec",
    "no", "vs", "est", "approx", "fig", "al", "etc", "e.g", "i.e",
}

_WS = re.compile(r"\s+")
_NON_ALPHA_DOT = re.compile(r"[^a-zA-Z.]")
_SINGLE_LOWER = re.compile(r"^[a-z]$")
_TRAILING_DOT = re.compile(r"\.$")


def _split_sentences(text: str) -> list[str]:
    """Naive but robust sentence splitter. Splits on ., !, ? followed by
    whitespace, while guarding common abbreviations and patterns."""
    normalized = _WS.sub(" ", text).strip()
    if not normalized:
        return []

    sentences: list[str] = []
    start = 0
    n = len(normalized)

    for i, ch in enumerate(normalized):
        if ch not in ".!?":
            continue

        # A boundary needs whitespace (or end-of-string) after the terminator.
        # "3.5" or "U.S.A" without a trailing space is not one.
        nxt = normalized[i + 1] if i + 1 < n else None
        if nxt is not None and nxt != " ":
            continue

        # Don't break after a known abbreviation or a single capital letter
        # (initials like "J. Powell").
        word = normalized[start:i].split(" ")[-1]
        word = _NON_ALPHA_DOT.sub("", word).lower()
        word = _TRAILING_DOT.sub("", word)
        if ch == "." and (word in ABBREVIATIONS or _SINGLE_LOWER.match(word)):
            continue

        sentence = normalized[start:i + 1].strip()
        if sentence:
            sentences.append(sentence)
        start = i + 1

    tail = normalized[start:].strip()
    if tail:
        sentences.append(tail)

    return _merge_fragments(sentences)


def _merge_fragments(sentences: list[str]) -> list[str]:
    """Glue too-short sentences onto the next one so every emitted chunk
    carries enough text to embed meaningfully. A trailing fragment with
    nothing after it is appended to the previous chunk instead."""
    out: list[str] = []
    buffer = ""
    for s in sentences:
        buffer = f"{buffer} {s}" if buffer else s
        if len(buffer) >= MIN_CHUNK_CHARS:
            out.append(buffer)
            buffer = ""
    if buffer:
        if out:
            out[-1] = f"{out[-1]} {buffer}"
        else:
            out.append(buffer)
    return out


def chunk_news(title: str, body: str) -> list[str]:
    """Article -> ordered chunks. The title is chunk 0 (highest-signal,
    almost always names the lead entity) followed by body sentences, capped.
    The title is dropped only if empty."""
    chunks: list[str] = []
    t = _WS.sub(" ", title).strip()
    if t:
        chunks.append(t)
    chunks.extend(_split_sentences(body)[:MAX_NEWS_CHUNKS])
    return _dedupe(chunks)


def _dedupe(chunks: list[str]) -> list[str]:
    """Order-preserving de-dupe (case-insensitive)."""
    seen: set[str] = set()
    out: list[str] = []
    for c in chunks:
        trimmed = c.strip()
        key = trimmed.lower()
        if not trimmed or key in seen:
            continue
        seen.add(key)
        out.append(trimmed)
    return out