llm-ingestion-okf/src/llm_ingestion_okf/bm25.py

"""BM25 over a bundle's concepts: the default ranking of `okf consume` (v1.1 C1).

WHY IT REPLACED THE FUSION. The earlier ranking fused three signals -- token
overlap with the title, token overlap with the body, and a document prior --
and two of them rewarded the wrong thing on a large collection: the document
prior favoured big documents full of common words, and the rarity weight gave
its MAXIMUM weight to a word that occurs nowhere in the collection. BM25 has
neither property, needs no new dependency, and ranks in milliseconds. The
measurement that chose it lives with the measurement; this module states the
mechanism.

TWO SIGNALS, FUSED BY RANK.

- **Passage** -- every body is cut into windows of `WINDOW_CHARS` characters
  every `WINDOW_STEP`, each window is a BM25 document, and a concept scores its
  BEST window (pure max). A narrow factual question is answered by one place
  in a concept, and a sum over windows -- even a damped one -- rewards a long
  concept for mentioning a word often.
- **Field** -- one BM25 document per concept: its title three times, its
  directory path and source file name twice, then its body. A broad question
  is answered by what a section is CALLED, and a passage never sees the title.

Fused by reciprocal rank (`RRF_K`, the constant the old fusion used). A signal
that scored a concept zero contributes nothing to it, and concepts that tie
within a signal share the group's first rank -- so neither a word the
collection lacks nor alphabetical order can lift a concept.

**A word the collection does not hold weighs exactly zero** (`idf` of a term
with `df == 0`), which is the property the old rarity weight had backwards.

Deterministic: every sort breaks ties by `concept_id`, and no float leaves this
module except as an ordering key.
"""

from __future__ import annotations

import functools
import math
import re
import unicodedata
from collections import Counter
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from llm_ingestion_okf.consume import Concept

#: BM25's two parameters, at the textbook values. Not tuned: the measurement
#: that chose this ranking used them as they are.
K1 = 1.2
B = 0.75

#: The passage window, in characters, and the step between window starts --
#: half a window, so every sentence sits whole inside at least one window.
WINDOW_CHARS = 500
WINDOW_STEP = 250

#: How many times the field signal repeats a concept's title, and its path.
TITLE_WEIGHT = 3
PATH_WEIGHT = 2

#: The rank-fusion constant, the same one the older fusion used.
RRF_K = 60

_TOKEN = re.compile(r"[0-9a-zà-öø-ÿ]+(?:[-_][0-9a-zà-öø-ÿ]+)*")

#: English and Norwegian function words. Norwegian because an operator asks in
#: Norwegian against a collection that may be English; a Norwegian function
#: word is then noise that could only ever match by accident. The last line is
#: the Norwegian ones spelled without their letters (`når` as `naar`), the way
#: ASCII-only text writes them, and the words that FRAME a question in either
#: language (`how often`, `hvor ofte`, `hva står i`, `what does it say`) -- they ask about a topic
#: without naming one, so read as topic words they would be "absent" from
#: every collection that never uses them (`coverage.absent_terms`).
STOPWORDS = frozenset(
    """a an and are as at be but by for if in into is it its of on or such that the
their then there these they this to was will with what which how when where who why
do does did can could should would i you your my me we our us not no yes from over
under about after before more most other some only own same so than too very s t just
don now am been being have has had having he she him her his hers them up out off
again further once here all any both each few nor
og i jeg det at en et den til er som på de med han av ikke der så var meg seg
men ett har om vi min mitt ha hadde hun nå over da ved fra du ut sin dem oss opp man
kan hans hvor eller hva skal selv her alle vil bli ble blitt kunne inn når være kom
noen noe ville dere hvordan gjør gjøre lar la flere
paa saa naa naar vaere gjoer gjoere
often many much ofte mange mye lenge hvilken hvilke hvilket hvorfor hvem bare også
ogsaa hver ingen uten mellom fordi før foer etter enn både baade denne dette disse
samme slik sånn saann står staar sier say says""".split()
)

_SUFFIXES = ("ingly", "edly", "ing", "ies", "ied", "ed", "es", "s", "ly", "er", "est")


@functools.lru_cache(maxsize=None)
def _stem(token: str) -> str:
    """A light English suffix strip, applied alike to question and text.

    Memoised because it is pure and a collection repeats its words: every
    question re-tokenises every window, so the same words are stemmed again
    and again.
    """
    if len(token) <= 3:
        return token
    for suffix in _SUFFIXES:
        if token.endswith(suffix) and len(token) - len(suffix) >= 3:
            return token[: -len(suffix)]
    return token


#: A number a document is known by -- `4.2`, `10.2-2`, `2.1.219` -- kept as ONE
#: token beside its pieces. Split on the dot, `4.2` is two single characters,
#: which are dropped, so a question naming a section by its number would match
#: nothing in the section it names. The shape is `consume`'s identifier rule.
_IDENTIFIER = re.compile(r"[0-9a-zà-öø-ÿ]*[0-9]+(?:[.-][0-9]+)+")
_DASH_TO_HYPHEN = str.maketrans(dict.fromkeys("‐‑‒–—―−", "-"))


def tokens(text: str) -> list[str]:
    """Casefolded, NFC, stopwords and single characters out, lightly stemmed,
    plus every identifier whole and unstemmed."""
    folded = unicodedata.normalize("NFC", text).casefold().translate(_DASH_TO_HYPHEN)
    words = [
        _stem(token)
        for token in _TOKEN.findall(folded)
        if len(token) > 1 and token not in STOPWORDS
    ]
    return words + [token for token in _IDENTIFIER.findall(folded) if token not in words]


class Index:
    """Okapi BM25 over a list of token lists."""

    def __init__(self, documents: Sequence[Sequence[str]]) -> None:
        self.size = len(documents)
        self.lengths = [len(document) for document in documents]
        self.average = sum(self.lengths) / self.size if self.size else 0.0
        self.postings: dict[str, list[tuple[int, int]]] = {}
        for position, document in enumerate(documents):
            for term, frequency in sorted(Counter(document).items()):
                self.postings.setdefault(term, []).append((position, frequency))

    def idf(self, terms: frozenset[str]) -> float:
        """Zero for terms no document holds: absence lifts nothing."""
        df = len(self.holders(terms))
        if df == 0:
            return 0.0
        return math.log(1 + (self.size - df + 0.5) / (df + 0.5))

    def holders(self, terms: frozenset[str]) -> set[int]:
        """Every document position holding at least one of `terms`."""
        return {position for term in terms for position, _ in self.postings.get(term, ())}

    def scores(self, query: Sequence[frozenset[str]]) -> dict[int, float]:
        """Positive scores only, keyed by document position.

        Each query element is a GROUP of terms read as one: its frequency in a
        document is the sum over the group, its `df` the documents holding any.
        A group of one is plain BM25.
        """
        out: dict[int, float] = {}
        for group in query:
            weight = self.idf(group)
            if weight == 0.0:
                continue
            frequencies: Counter[int] = Counter()
            for term in sorted(group):
                for position, frequency in self.postings.get(term, ()):
                    frequencies[position] += frequency
            for position, frequency in sorted(frequencies.items()):
                norm = K1 * (1 - B + B * self.lengths[position] / self.average)
                out[position] = out.get(position, 0.0) + weight * (
                    frequency * (K1 + 1) / (frequency + norm)
                )
        return out


def windows(body: str) -> list[tuple[int, str]]:
    """`(start, text)` for every window of a body; one empty-start window if blank."""
    if not body.strip():
        return []
    out = []
    for start in range(0, len(body), WINDOW_STEP):
        chunk = body[start : start + WINDOW_CHARS]
        if chunk.strip():
            out.append((start, chunk))
        if start + WINDOW_CHARS >= len(body):
            break
    return out


def field_text(concept: Concept, body: str, *, shared: int = 0, own_source: bool = True) -> str:
    """Title and path weighted up, then the body: the field signal's document.

    `shared` is how many leading id segments EVERY concept carries, and
    `own_source` whether the source file name separates this concept from any
    other. What every concept carries separates nothing, and a question naming
    it would match them all -- the defect `consume.shared_id_prefix` closed for
    the older ranking, and the reason both are dropped here.
    """
    path = " ".join(concept.concept_id.split("/")[shared:]).replace("-", " ")
    source = concept.source_file.removesuffix(".md").replace("-", " ") if own_source else ""
    return f"{concept.title} " * TITLE_WEIGHT + f"{path} {source} " * PATH_WEIGHT + body


def _shared_segments(concept_ids: Sequence[str]) -> int:
    """How many leading DIRECTORY segments every id shares (never the leaf)."""
    if not concept_ids:
        return 0
    split = [concept_id.split("/")[:-1] for concept_id in concept_ids]
    count = 0
    for segments in zip(*split):
        if len(set(segments)) != 1:
            break
        count += 1
    return count


def query_groups(query: Sequence[str], vocabulary: frozenset[str]) -> list[frozenset[str]]:
    """Each question term as the group of collection terms it is read as.

    A term the collection holds is read as itself and nothing else. A term it
    does NOT hold -- an inflection, a compound, a word in another language --
    is read as the collection's words it shares a leading WORD with, by
    `consume.tokens_match`, the rule the older ranking measured for Norwegian
    inflection and compounding. So an absent word lifts nothing by itself: it
    reaches the collection only through a relative the collection uses, at that
    relative's `idf` and never at a weight of its own.
    """
    from llm_ingestion_okf.consume import MIN_SHARED_PREFIX, tokens_match

    by_prefix: dict[str, list[str]] = {}
    for term in sorted(vocabulary):
        by_prefix.setdefault(term[:MIN_SHARED_PREFIX], []).append(term)
    groups = []
    for term in dict.fromkeys(query):
        if term in vocabulary:
            groups.append(frozenset({term}))
            continue
        groups.append(
            frozenset(
                candidate
                for candidate in by_prefix.get(term[:MIN_SHARED_PREFIX], ())
                if tokens_match(term, candidate, stems=vocabulary)
            )
        )
    return groups


@dataclass(frozen=True)
class Ranking:
    """Every concept best first, as `(concept, fused, lexical)`, plus where
    each concept's best window starts in the body that was searched."""

    ranked: list[tuple[Concept, float, int]]
    best_window: dict[str, int]
    #: The question's words the collection holds in NO form, in question order.
    absent: tuple[str, ...] = ()


def _fuse(fused: dict[str, float], scores: dict[str, float]) -> None:
    """Add one signal's reciprocal ranks; equal scores share the first rank."""
    order = sorted(scores, key=lambda key: (-scores[key], key))
    start = 0
    while start < len(order):
        stop = start
        while stop < len(order) and scores[order[stop]] == scores[order[start]]:
            stop += 1
        contribution = 1.0 / (RRF_K + start + 1)
        for concept_id in order[start:stop]:
            fused[concept_id] += contribution
        start = stop


@dataclass(frozen=True)
class Prepared:
    """Everything `rank` reads that does not depend on the question.

    Built once per load of a bundle, so a call asking several sub-questions
    (`consume.build_multi_payload`) tokenises and indexes the collection once
    and ranks it once per sub-question. `rank` builds one itself when not
    given one, so a single question pays exactly what it always paid.
    """

    concepts: tuple[Concept, ...]
    field_documents: tuple[tuple[str, ...], ...]
    field_index: Index
    vocabulary: frozenset[str]
    owners: tuple[int, ...]
    starts: tuple[int, ...]
    passage_index: Index


def prepare(concepts: Sequence[Concept], *, bodies: Sequence[str] | None = None) -> Prepared:
    """Index `concepts` for ranking: the field documents and the passages.

    `bodies` is the text searched per concept (defaults to each `body`); the
    caller passes the body without the door's link line, so what is searched
    is what the older ranking searched.
    """
    texts = list(bodies) if bodies is not None else [concept.body for concept in concepts]
    shared = _shared_segments([concept.concept_id for concept in concepts])
    own_source = len({concept.source_file for concept in concepts}) > 1
    field_documents = tuple(
        tuple(tokens(field_text(concept, text, shared=shared, own_source=own_source)))
        for concept, text in zip(concepts, texts, strict=True)
    )
    field_index = Index(field_documents)

    owners: list[int] = []
    starts: list[int] = []
    passages: list[list[str]] = []
    for position, text in enumerate(texts):
        cut = windows(text) or [(0, concepts[position].title)]
        for start, chunk in cut:
            owners.append(position)
            starts.append(start)
            passages.append(tokens(chunk))
    return Prepared(
        concepts=tuple(concepts),
        field_documents=field_documents,
        field_index=field_index,
        vocabulary=frozenset(field_index.postings),
        owners=tuple(owners),
        starts=tuple(starts),
        passage_index=Index(passages),
    )


def rank(
    concepts: Sequence[Concept],
    question: str,
    *,
    bodies: Sequence[str] | None = None,
    prepared: Prepared | None = None,
) -> Ranking:
    """Rank `concepts` for `question`.

    `bodies` is the text searched per concept (defaults to each `body`); the
    caller passes the body without the door's link line, so what is searched
    is what the older ranking searched. `prepared` is `prepare`'s result for
    the same `concepts` and `bodies`, given when one load answers several
    questions; the ranking is the same either way.
    """
    if prepared is None:
        prepared = prepare(concepts, bodies=bodies)
    concepts = prepared.concepts
    query = tokens(question)
    groups = query_groups(query, prepared.vocabulary)
    field = {
        concepts[position].concept_id: score
        for position, score in prepared.field_index.scores(groups).items()
    }

    passage: dict[str, float] = {}
    best_window: dict[str, int] = {}
    for window, score in sorted(prepared.passage_index.scores(groups).items()):
        concept_id = concepts[prepared.owners[window]].concept_id
        if score > passage.get(concept_id, 0.0):
            passage[concept_id] = score
            best_window[concept_id] = prepared.starts[window]

    fused = {concept.concept_id: 0.0 for concept in concepts}
    _fuse(fused, passage)
    _fuse(fused, field)

    asked = [group for group in groups if prepared.field_index.idf(group) > 0.0]
    lexical = {
        concept.concept_id: sum(1 for group in asked if group & held)
        for concept, held in (
            (concept, set(document))
            for concept, document in zip(concepts, prepared.field_documents, strict=True)
        )
    }
    by_id = {concept.concept_id: concept for concept in concepts}
    order = sorted(fused, key=lambda key: (-fused[key], key))
    return Ranking(
        ranked=[(by_id[key], fused[key], lexical[key]) for key in order],
        best_window=best_window,
        absent=_absent(query, groups),
    )


def _absent(query: Sequence[str], groups: Sequence[frozenset[str]]) -> tuple[str, ...]:
    return tuple(
        term for term, group in zip(dict.fromkeys(query), groups, strict=True) if not group
    )


def absent_terms(
    concepts: Sequence[Concept], question: str, *, bodies: Sequence[str]
) -> tuple[str, ...]:
    """The question's words the collection holds in no form -- not as written
    and not through a relative (`query_groups`). The same reading `rank`
    reports, for a caller ranking some other way."""
    shared = _shared_segments([concept.concept_id for concept in concepts])
    own_source = len({concept.source_file for concept in concepts}) > 1
    vocabulary = frozenset(
        term
        for concept, text in zip(concepts, bodies, strict=True)
        for term in tokens(field_text(concept, text, shared=shared, own_source=own_source))
    )
    query = tokens(question)
    return _absent(query, query_groups(query, vocabulary))