Модуль 17. Эмбеддинги, поиск и RAG

import importlib.util
import subprocess
import sys


def ensure_package(import_name: str, pip_name: str | None = None) -> None:
    """Install one missing package without touching an already prepared runtime."""

    if importlib.util.find_spec(import_name) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pip_name or import_name])


ensure_package("numpy")

import hashlib
import json
import math
import re
from collections import Counter
from dataclasses import dataclass
from typing import Mapping, Sequence

import numpy as np

np.set_printoptions(precision=5, suppress=True)
print('Окружение готово.')

import json, math, re, hashlib
from collections import Counter
from dataclasses import dataclass, asdict
from typing import Mapping, Sequence
import numpy as np
WORD_RE = re.compile(r"[^\W_]+", re.UNICODE)
REFERENCE_JSON = '{"ann": {"ann_flops": 4096000.0, "ann_milliseconds": 0.40959999999999996, "exact_flops": 20480000000.0, "exact_seconds": 2.048, "speedup": 5000.0}, "cascade": {"ann_ms": 0.40959999999999996, "context_tokens": 3712, "decode_ms": 1432.6647564469915, "decode_share": 0.8639290480052048, "embed_ms": 0.07757575757575758, "prefill_ms": 120.43377777777778, "rerank_ms": 104.72727272727272, "retrieval_before_prefill_share": 0.06344667718450364, "total_ms": 1658.3129827096177}, "chunk_example": [{"chunk_id": "demo:c000", "end_char": 53, "end_token": 8, "source_id": "demo", "start_char": 0, "start_token": 0, "text": "retrieval first splits a long source into overlapping", "token_count": 8}, {"chunk_id": "demo:c001", "end_char": 94, "end_token": 13, "source_id": "demo", "start_char": 30, "start_token": 5, "text": "source into overlapping windows while keeping source identifiers", "token_count": 8}, {"chunk_id": "demo:c002", "end_char": 125, "end_token": 18, "source_id": "demo", "start_char": 68, "start_token": 10, "text": "keeping source identifiers offsets and citations so later", "token_count": 8}, {"chunk_id": "demo:c003", "end_char": 159, "end_token": 22, "source_id": "demo", "start_char": 107, "start_token": 15, "text": "citations so later context packing remains traceable", "token_count": 7}], "corpus": {"binary_mib": 7.26611328125, "chunk_size": 400, "chunks": 59524, "encoder_seconds_floor": 57.72024242424242, "fp32_mib": 232.515625, "int8_mib": 58.12890625, "overlap": 64, "total_tokens": 20000000}, "datasets": {"ru": {"documents": [{"document_id": "d00", "source_id": "ann", "text": "HNSW строит многослойный граф малого мира для приближённого поиска ближайших соседей среди векторов."}, {"document_id": "d01", "source_id": "ann", "text": "IVF PQ делит векторы на кластеры и кодирует остатки, уменьшая память индекса для сотен миллионов объектов."}, {"document_id": "d02", "source_id": "hybrid", "text": "Reciprocal rank fusion RRF объединяет ранги BM25 и dense поиска без калибровки их оценок."}, {"document_id": "d03", "source_id": "sparse", "text": "BM25 использует частоту терма, обратную документную частоту и нормировку длины документа."}, {"document_id": "d04", "source_id": "dense", "text": "Dense bi encoder независимо кодирует запрос и документ в векторы, после чего сравнивает их косинусом."}, {"document_id": "d05", "source_id": "rerank", "text": "Cross encoder совместно читает запрос и документ и точнее оценивает релевантность, но требует отдельного прохода на каждую пару."}, {"document_id": "d06", "source_id": "chunking", "text": "Overlap между соседними чанками сохраняет локальный контекст на границе разбиения документа."}, {"document_id": "d07", "source_id": "chunking", "text": "Late chunking сначала кодирует длинный документ целиком, поэтому местоимения и анафоры сохраняют контекст после нарезки."}, {"document_id": "d08", "source_id": "routing", "text": "Long context помещает целый длинный документ в промпт и подходит для целостного анализа одного файла."}, {"document_id": "d09", "source_id": "rag", "text": "RAG извлекает релевантные фрагменты корпуса, собирает контекст с цитатами и передаёт его генератору."}, {"document_id": "d10", "source_id": "metrics", "text": "nDCG учитывает градуированную релевантность и логарифмически штрафует низкие позиции."}, {"document_id": "d11", "source_id": "metrics", "text": "MRR измеряет обратный ранг первого релевантного результата и игнорирует остальные найденные документы."}, {"document_id": "d12", "source_id": "packing", "text": "Context packing удаляет дубликаты, соблюдает токенный бюджет и сохраняет ссылки на источники."}], "queries": [{"query_id": "q0", "relevance": {"d02": 3, "d09": 1}, "text": "как слить результаты bm25 и dense"}, {"query_id": "q1", "relevance": {"d00": 3, "d01": 1}, "text": "граф соседей для approximate nearest neighbor"}, {"query_id": "q2", "relevance": {"d10": 1, "d11": 3}, "text": "позиция первого полезного результата"}, {"query_id": "q3", "relevance": {"d06": 2, "d07": 3}, "text": "как сохранить анафоры после нарезки документа"}, {"query_id": "q4", "relevance": {"d04": 1, "d05": 3}, "text": "почему cross encoder нельзя запустить на всём корпусе"}, {"query_id": "q5", "relevance": {"d08": 3, "d09": 2}, "text": "когда лучше прочитать весь длинный файл вместо retrieval"}, {"query_id": "q6", "relevance": {"d09": 2, "d12": 3}, "text": "context packing: как собрать найденные фрагменты с цитатами в лимит токенов"}]}}, "long_context_vs_rag": {"kv_bytes_per_token": 131072, "long_kv_gib": 24.4140625, "long_prefill_seconds": 6.488888888888889, "long_tokens": 200000, "rag_kv_gib": 0.453125, "rag_prefill_seconds": 0.12043377777777778, "rag_tokens": 3712, "ratio": 53.87931034482759}, "retrieval_results": {"ru": {"context_pack_q6": {"citations": ["[rag:d09]", "[packing:d12]", "[chunking:d06]"], "max_tokens": 42, "rendered_text": "[rag:d09] RAG извлекает релевантные фрагменты корпуса, собирает контекст с цитатами и передаёт его генератору.\\n\\n[packing:d12] Context packing удаляет дубликаты, соблюдает токенный бюджет и сохраняет ссылки на источники.\\n\\n[chunking:d06] Overlap между соседними чанками сохраняет локальный контекст на границе разбиения документа.", "selected_ids": ["d09", "d12", "d06"], "skipped": {"d00": "budget", "d04": "budget", "d07": "source_limit", "d08": "budget", "d11": "budget"}, "used_tokens": 42}, "metrics": {"bm25": {"mrr": 0.9285714285714286, "ndcg_at_5": 0.8719917335378166, "recall_at_5": 0.7857142857142857}, "dense_lsa": {"mrr": 0.7619047619047619, "ndcg_at_5": 0.7552817156182359, "recall_at_5": 0.8571428571428571}, "reranked": {"mrr": 1.0, "ndcg_at_5": 0.927756548206894, "recall_at_5": 0.8571428571428571}, "rrf": {"mrr": 0.9285714285714286, "ndcg_at_5": 0.8660296227599356, "recall_at_5": 0.7857142857142857}}, "rankings": {"bm25": {"q0": ["d02", "d03", "d04", "d05", "d10", "d11", "d12", "d09", "d07", "d08", "d01", "d00", "d06"], "q1": ["d00", "d08", "d01", "d02", "d03", "d04", "d05", "d06", "d07", "d09", "d10", "d11", "d12"], "q2": ["d11", "d00", "d01", "d02", "d03", "d04", "d05", "d06", "d07", "d08", "d09", "d10", "d12"], "q3": ["d07", "d03", "d06", "d04", "d00", "d01", "d02", "d05", "d08", "d09", "d10", "d11", "d12"], "q4": ["d05", "d04", "d06", "d12", "d01", "d00", "d02", "d03", "d07", "d08", "d09", "d10", "d11"], "q5": ["d07", "d08", "d00", "d01", "d02", "d03", "d04", "d05", "d06", "d09", "d10", "d11", "d12"], "q6": ["d09", "d12", "d08", "d11", "d04", "d00", "d01", "d02", "d03", "d05", "d06", "d07", "d10"]}, "dense_lsa": {"q0": ["d03", "d02", "d11", "d04", "d00", "d10", "d08", "d05", "d07", "d09", "d12", "d06", "d01"], "q1": ["d00", "d01", "d02", "d03", "d04", "d08", "d12", "d11", "d06", "d07", "d10", "d05", "d09"], "q2": ["d11", "d04", "d03", "d10", "d08", "d07", "d02", "d05", "d09", "d06", "d12", "d00", "d01"], "q3": ["d07", "d09", "d08", "d06", "d11", "d04", "d12", "d03", "d02", "d10", "d05", "d01", "d00"], "q4": ["d05", "d10", "d12", "d11", "d04", "d06", "d01", "d03", "d02", "d08", "d07", "d00", "d09"], "q5": ["d07", "d08", "d09", "d06", "d04", "d11", "d12", "d03", "d01", "d02", "d10", "d05", "d00"], "q6": ["d07", "d08", "d09", "d06", "d04", "d11", "d12", "d03", "d10", "d01", "d02", "d05", "d00"]}, "reranked": {"q0": ["d02", "d03", "d04", "d11", "d05", "d10", "d08", "d00", "d12", "d09"], "q1": ["d00", "d01", "d08", "d03", "d05", "d12", "d02", "d04", "d06", "d11"], "q2": ["d11", "d10", "d00", "d05", "d04", "d02", "d03", "d01", "d08", "d07", "d06"], "q3": ["d07", "d04", "d06", "d03", "d08", "d11", "d09", "d00", "d01", "d02", "d12", "d05"], "q4": ["d05", "d04", "d01", "d06", "d12", "d10", "d11", "d03", "d00", "d02"], "q5": ["d08", "d07", "d03", "d01", "d00", "d04", "d06", "d09", "d02", "d11", "d12", "d05"], "q6": ["d09", "d12", "d08", "d11", "d04", "d06", "d07", "d00", "d01", "d02", "d03"]}, "rrf": {"q0": ["d02", "d03", "d04", "d11", "d10", "d05", "d00", "d08", "d12", "d09"], "q1": ["d00", "d01", "d02", "d08", "d03", "d04", "d05", "d12", "d06", "d11"], "q2": ["d11", "d04", "d03", "d02", "d05", "d00", "d01", "d10", "d08", "d07", "d06"], "q3": ["d07", "d06", "d03", "d04", "d09", "d08", "d00", "d11", "d01", "d02", "d12", "d05"], "q4": ["d05", "d04", "d12", "d06", "d01", "d03", "d10", "d11", "d00", "d02"], "q5": ["d07", "d08", "d04", "d03", "d00", "d09", "d01", "d06", "d02", "d11", "d12", "d05"], "q6": ["d09", "d08", "d12", "d11", "d04", "d07", "d06", "d00", "d01", "d02", "d03"]}}, "rrf_trace_q0": [{"contributions": {"bm25": 0.01639344262295082, "dense_lsa": 0.016129032258064516}, "document_id": "d02", "rank": 1, "score": 0.03252247488101534, "source_ranks": {"bm25": 1, "dense_lsa": 2}}, {"contributions": {"bm25": 0.016129032258064516, "dense_lsa": 0.01639344262295082}, "document_id": "d03", "rank": 2, "score": 0.03252247488101534, "source_ranks": {"bm25": 2, "dense_lsa": 1}}, {"contributions": {"bm25": 0.015873015873015872, "dense_lsa": 0.015625}, "document_id": "d04", "rank": 3, "score": 0.03149801587301587, "source_ranks": {"bm25": 3, "dense_lsa": 4}}, {"contributions": {"bm25": 0.015151515151515152, "dense_lsa": 0.015873015873015872}, "document_id": "d11", "rank": 4, "score": 0.031024531024531024, "source_ranks": {"bm25": 6, "dense_lsa": 3}}, {"contributions": {"bm25": 0.015384615384615385, "dense_lsa": 0.015151515151515152}, "document_id": "d10", "rank": 5, "score": 0.030536130536130537, "source_ranks": {"bm25": 5, "dense_lsa": 6}}]}}, "schema_version": 2, "storage": {"binary_gib": 1.1920928955078125, "dimension": 1024, "fp16_gib": 19.073486328125, "fp32_gib": 38.14697265625, "int8_gib": 9.5367431640625, "mrl256_int8_gib": 2.384185791015625, "num_vectors": 10000000}}'
assert hashlib.sha256(REFERENCE_JSON.encode()).hexdigest() == 'ef50d75622e2e80154a8400d10f7eb93a914299ade6d89f634440414be674286'
REFERENCE = json.loads(REFERENCE_JSON)
LANG = 'ru'
DATASETS = REFERENCE["datasets"]
data = DATASETS[LANG]
ids = [row["document_id"] for row in data["documents"]]
texts = [row["text"] for row in data["documents"]]
documents = {row["document_id"]: row for row in data["documents"]}
queries = data["queries"]
stored = {}

def _positive_int(name: str, value: int) -> int:
    """Validate and return one strictly positive integer."""

    if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)

def _positive_finite(name: str, value: float) -> float:
    """Validate and return one strictly positive finite float."""

    value = float(value)
    if not math.isfinite(value) or value <= 0.0:
        raise ValueError(f"{name} must be finite and positive")
    return value

def tokenize(text: str) -> tuple[str, ...]:
    """Find word spans in the original text, then casefold each word for search."""

    if not isinstance(text, str):
        raise TypeError("text must be a string")
    return tuple(match.group().casefold() for match in WORD_RE.finditer(text))

@dataclass(frozen=True)
class Chunk:
    """One contiguous token window with source and offset metadata."""

    chunk_id: str
    source_id: str
    start_token: int
    end_token: int
    start_char: int
    end_char: int
    text: str
    token_count: int

def chunk_document(
    source_id: str,
    text: str,
    *,
    chunk_size: int,
    overlap: int,
) -> tuple[Chunk, ...]:
    """Split one document into deterministic overlapping token windows.

    Whitespace-only documents return an empty tuple explicitly.  Non-empty documents
    are covered without gaps; the final chunk is not duplicated when the last regular
    window already reaches the end.
    """

    if not isinstance(source_id, str) or not source_id.strip():
        raise ValueError("source_id must be a non-empty string")
    chunk_size = _positive_int("chunk_size", chunk_size)
    if not isinstance(overlap, (int, np.integer)) or isinstance(overlap, bool):
        raise TypeError("overlap must be an integer")
    overlap = int(overlap)
    if overlap < 0 or overlap >= chunk_size:
        raise ValueError("overlap must satisfy 0 <= overlap < chunk_size")
    if not isinstance(text, str):
        raise TypeError("text must be a string")
    spans = tuple(WORD_RE.finditer(text))
    if not spans:
        return ()
    stride = chunk_size - overlap
    chunks: list[Chunk] = []
    start = 0
    index = 0
    while start < len(spans):
        end = min(start + chunk_size, len(spans))
        # Character offsets refer to the original Python string, not bytes or
        # casefolded text. Preserve case, punctuation, and spacing inside the span.
        char_start = spans[start].start()
        char_end = spans[end - 1].end()
        chunks.append(
            Chunk(
                chunk_id=f"{source_id}:c{index:03d}",
                source_id=source_id,
                start_token=start,
                end_token=end,
                start_char=char_start,
                end_char=char_end,
                text=text[char_start:char_end],
                token_count=end - start,
            )
        )
        if end == len(spans):
            break
        start += stride
        index += 1
    return tuple(chunks)

def analytical_chunk_count(total_tokens: int, chunk_size: int, overlap: int) -> int:
    """Return the exact sliding-window count for a tokenized document."""

    total_tokens = _positive_int("total_tokens", total_tokens)
    chunk_size = _positive_int("chunk_size", chunk_size)
    if not isinstance(overlap, (int, np.integer)) or isinstance(overlap, bool):
        raise TypeError("overlap must be an integer")
    overlap = int(overlap)
    if overlap < 0 or overlap >= chunk_size:
        raise ValueError("overlap must satisfy 0 <= overlap < chunk_size")
    if total_tokens <= chunk_size:
        return 1
    stride = chunk_size - overlap
    return (total_tokens - chunk_size + stride - 1) // stride + 1

def _validate_ids(ids):
    """IDs remain nonempty strings; coercion must not merge 1 and '1'."""
    if any(not isinstance(i, str) or not i.strip() for i in ids):
        raise ValueError("document IDs must be nonempty strings")
    if len(ids) != len(set(ids)):
        raise ValueError("document IDs must be unique")

def _validate_judgments(ranking, relevance):
    """Reject duplicate ranks, fractional/negative grades and empty gold sets."""
    _validate_ids(ranking)
    _validate_ids(list(relevance))
    if any(isinstance(g, bool) or not isinstance(g, (int, np.integer))
           or not 0 <= g <= 30 for g in relevance.values()):
        raise ValueError("grades must be integers from 0 through 30")
    if not any(g > 0 for g in relevance.values()):
        raise ValueError("at least one positive judgment is required")

def stable_ranking(scores: Sequence[float], document_ids: Sequence[str]) -> tuple[str, ...]:
    """Sort scores descending with document ID as a deterministic tie-break."""

    score_array = np.asarray(scores, dtype=np.float64)
    if score_array.ndim != 1 or score_array.size != len(document_ids):
        raise ValueError("scores and document_ids must be aligned rank-1 sequences")
    if not np.all(np.isfinite(score_array)):
        raise ValueError("scores contain NaN or infinity")
    _validate_ids(document_ids)
    if len(set(document_ids)) != len(document_ids):
        raise ValueError("document_ids must be unique")
    return tuple(
        document_id
        for _, document_id in sorted(
            zip(score_array.tolist(), map(str, document_ids)),
            key=lambda pair: (-pair[0], pair[1]),
        )
    )

@dataclass(frozen=True)
class SearchResult:
    """One ranked document and its branch score."""

    document_id: str
    score: float
    rank: int
sample = "one two three four five six seven eight nine ten"
chunks = chunk_document("sample", sample, chunk_size=4, overlap=1)
assert len(chunks) == analytical_chunk_count(10, 4, 1) == 3
assert set(i for c in chunks for i in range(c.start_token,c.end_token)) == set(range(10))
print([(c.start_token,c.end_token,c.text) for c in chunks])
original = "TS-999:  UPDATE\nfailed — İstanbul; Straße."
quotes = chunk_document("case", original, chunk_size=4, overlap=1)
assert all(c.text == original[c.start_char:c.end_char] for c in quotes)
assert quotes[0].text == "TS-999:  UPDATE\nfailed"
assert tokenize(quotes[-1].text) == tokenize(original)[quotes[-1].start_token:quotes[-1].end_token]

class BM25Index:
    """A small exact BM25 index with explicit empty-input policy."""

    def __init__(
        self,
        document_ids: Sequence[str],
        texts: Sequence[str],
        *,
        k1: float = 1.5,
        b: float = 0.75,
    ) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        self.tokens = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in self.tokens):
            raise ValueError("empty documents must be removed before indexing")
        self.k1 = _positive_finite("k1", k1)
        self.b = float(b)
        if not math.isfinite(self.b) or not 0.0 <= self.b <= 1.0:
            raise ValueError("b must lie in [0,1]")
        n_documents = len(self.tokens)
        self.average_length = sum(map(len, self.tokens)) / n_documents
        self.document_frequency = Counter(
            token for document_tokens in self.tokens for token in set(document_tokens)
        )
        self.idf = {
            token: math.log(
                (n_documents - frequency + 0.5) / (frequency + 0.5) + 1.0
            )
            for token, frequency in self.document_frequency.items()
        }

    def scores(self, query: str) -> np.ndarray:
        """Return one BM25 score per indexed document."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        if not any(token in self.idf for token in query_tokens):
            raise ValueError("query has no indexed lexical terms")
        scores = np.zeros(len(self.tokens), dtype=np.float64)
        for row, document_tokens in enumerate(self.tokens):
            frequencies = Counter(document_tokens)
            length_norm = 1.0 - self.b + self.b * len(document_tokens) / self.average_length
            for token in query_tokens:
                tf = frequencies.get(token, 0)
                if tf == 0 or token not in self.idf:
                    continue
                scores[row] += (
                    self.idf[token]
                    * tf
                    * (self.k1 + 1.0)
                    / (tf + self.k1 * length_norm)
                )
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank all documents or the requested top-k with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )
bm25 = BM25Index(ids, texts)
for q in queries:
    stored[q["query_id"]] = {"bm25": bm25.search(q["text"])}
assert [r.document_id for r in stored["q0"]["bm25"]] == REFERENCE["retrieval_results"][LANG]["rankings"]["bm25"]["q0"]
print(stored["q0"]["bm25"][:3])

class LSAIndex:
    """Exact cosine search over dense TF-IDF + SVD document vectors."""

    def __init__(self, document_ids: Sequence[str], texts: Sequence[str], *, dimension: int = 4) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        tokenized = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in tokenized):
            raise ValueError("empty documents must be removed before indexing")
        vocabulary = sorted({token for tokens in tokenized for token in tokens})
        self.vocabulary = {token: index for index, token in enumerate(vocabulary)}
        counts = np.zeros((len(tokenized), len(vocabulary)), dtype=np.float64)
        document_frequency = np.zeros(len(vocabulary), dtype=np.float64)
        for row, tokens in enumerate(tokenized):
            frequencies = Counter(tokens)
            for token, count in frequencies.items():
                counts[row, self.vocabulary[token]] = count
                document_frequency[self.vocabulary[token]] += 1.0
        self.idf = np.log((1.0 + len(tokenized)) / (1.0 + document_frequency)) + 1.0
        tf = np.where(counts > 0.0, 1.0 + np.log(np.maximum(counts, 1.0)), 0.0)
        tfidf = tf * self.idf
        _, singular_values, vt = np.linalg.svd(tfidf, full_matrices=False)
        numerical_rank = int(np.sum(singular_values > 1e-10))
        dimension = min(_positive_int("dimension", dimension), numerical_rank)
        self.projection = vt[:dimension].T
        dense = tfidf @ self.projection
        norms = np.linalg.norm(dense, axis=1, keepdims=True)
        if np.any(norms <= 0.0):
            raise ValueError("at least one document has a zero dense representation")
        self.document_vectors = dense / norms

    def encode_query(self, query: str) -> np.ndarray:
        """Map one non-empty in-vocabulary query to a unit dense vector."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        counts = Counter(query_tokens)
        sparse = np.zeros(len(self.vocabulary), dtype=np.float64)
        for token, count in counts.items():
            position = self.vocabulary.get(token)
            if position is not None:
                sparse[position] = (1.0 + math.log(count)) * self.idf[position]
        dense = sparse @ self.projection
        norm = float(np.linalg.norm(dense))
        if norm <= 0.0:
            raise ValueError("query contains no terms represented by the dense index")
        return dense / norm

    def scores(self, query: str) -> np.ndarray:
        """Return exact cosine scores against every document vector."""

        scores = self.document_vectors @ self.encode_query(query)
        if not np.all(np.isfinite(scores)):
            raise FloatingPointError("dense search produced NaN or infinity")
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank exact dense cosine scores with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )
dense = LSAIndex(ids, texts, dimension=4)
assert dense.document_vectors.shape == (13,4)
assert np.allclose(np.linalg.norm(dense.document_vectors,axis=1),1)
for q in queries:
    stored[q["query_id"]]["dense_lsa"] = dense.search(q["text"])
print(dense.document_vectors.shape)

@dataclass(frozen=True)
class FusedResult:
    """One RRF result with a per-branch contribution trace."""

    document_id: str
    score: float
    rank: int
    contributions: Mapping[str, float]
    source_ranks: Mapping[str, int]

def reciprocal_rank_fusion(
    rankings: Mapping[str, Sequence[str]],
    *,
    k: float = 60.0,
    weights: Mapping[str, float] | None = None,
    depth: int | None = None,
) -> tuple[FusedResult, ...]:
    """Accumulate ranks in linear time, then sort unique candidates.

    Ranks are one-based. Duplicate document IDs inside one branch are rejected because
    they would double-count evidence.  The implementation builds rank maps once and
    never calls ``list.index`` inside a loop.
    """

    k = _positive_finite("k", k)
    if not rankings:
        raise ValueError("at least one ranking is required")
    if depth is not None:
        depth = _positive_int("depth", depth)
    branch_weights = {name: 1.0 for name in rankings}
    if weights is not None:
        if set(weights) != set(rankings):
            raise ValueError("weights must contain exactly the ranking branch names")
        branch_weights = {
            name: _positive_finite(f"weight[{name}]", weights[name]) for name in rankings
        }
    scores: dict[str, float] = {}
    contributions: dict[str, dict[str, float]] = {}
    source_ranks: dict[str, dict[str, int]] = {}
    for branch_name, ranking in rankings.items():
        _validate_ids(ranking)
        sequence = tuple(ranking)
        if len(sequence) != len(set(sequence)):
            raise ValueError(f"ranking {branch_name!r} contains duplicate document IDs")
        if depth is not None:
            sequence = sequence[:depth]
        for rank, document_id in enumerate(sequence, start=1):
            contribution = branch_weights[branch_name] / (k + rank)
            scores[document_id] = scores.get(document_id, 0.0) + contribution
            contributions.setdefault(document_id, {})[branch_name] = contribution
            source_ranks.setdefault(document_id, {})[branch_name] = rank
    ordered = sorted(scores, key=lambda document_id: (-scores[document_id], document_id))
    return tuple(
        FusedResult(
            document_id=document_id,
            score=float(scores[document_id]),
            rank=rank,
            contributions=dict(sorted(contributions[document_id].items())),
            source_ranks=dict(sorted(source_ranks[document_id].items())),
        )
        for rank, document_id in enumerate(ordered, start=1)
    )
for q in queries:
    row = stored[q["query_id"]]
    row["rrf"] = reciprocal_rank_fusion(
        {name:[r.document_id for r in row[name]] for name in ("bm25","dense_lsa")},
        depth=8)
print([asdict(r) for r in stored["q0"]["rrf"][:3]])

def _character_trigrams(token: str) -> frozenset[str]:
    """Return boundary-aware character trigrams for one token."""

    wrapped = f"^{token}$"
    return frozenset(wrapped[index : index + 3] for index in range(max(0, len(wrapped) - 2)))

def _token_similarity(left: str, right: str) -> float:
    """Return Jaccard similarity between two token trigram sets."""

    a = _character_trigrams(left)
    b = _character_trigrams(right)
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)

def token_interaction_score(query: str, document: str) -> float:
    """Score a query/document pair with a transparent MaxSim-style proxy.

    This is not a neural cross-encoder.  It deliberately exposes the per-query-token
    maximum interaction used by late-interaction retrieval families.
    """

    query_tokens = tokenize(query)
    document_tokens = tokenize(document)
    if not query_tokens:
        raise ValueError("query must contain at least one word token")
    if not document_tokens:
        raise ValueError("document must contain at least one word token")
    return float(
        sum(
            max(_token_similarity(query_token, document_token) for document_token in document_tokens)
            for query_token in query_tokens
        )
        / len(query_tokens)
    )

def rerank_candidates(
    query: str,
    candidate_ids: Sequence[str],
    documents: Mapping[str, str],
) -> tuple[SearchResult, ...]:
    """Rerank a candidate list with the token-interaction proxy."""

    _validate_ids(candidate_ids)
    candidate_ids = tuple(candidate_ids)
    if not candidate_ids:
        raise ValueError("candidate_ids must be non-empty")
    if len(candidate_ids) != len(set(candidate_ids)):
        raise ValueError("candidate_ids must be unique")
    missing = [document_id for document_id in candidate_ids if document_id not in documents]
    if missing:
        raise KeyError(f"missing documents: {missing}")
    scores = [token_interaction_score(query, documents[document_id]) for document_id in candidate_ids]
    ranking = stable_ranking(scores, candidate_ids)
    score_by_id = dict(zip(candidate_ids, scores))
    return tuple(
        SearchResult(document_id, float(score_by_id[document_id]), rank)
        for rank, document_id in enumerate(ranking, start=1)
    )
for q in queries:
    row = stored[q["query_id"]]
    row["reranked"] = rerank_candidates(q["text"], [r.document_id for r in row["rrf"][:8]],
                                       dict(zip(ids,texts)))
print(stored["q0"]["reranked"][:3])

def reciprocal_rank(ranking: Sequence[str], relevance: Mapping[str, int]) -> float:
    """Return reciprocal rank of the first positively relevant document."""
    _validate_judgments(ranking, relevance)

    if not any(int(grade) > 0 for grade in relevance.values()):
        raise ValueError("relevance must contain at least one positive grade")
    for rank, document_id in enumerate(ranking, start=1):
        if int(relevance.get(document_id, 0)) > 0:
            return 1.0 / rank
    return 0.0

def recall_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return set recall among all positively relevant documents."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    relevant = {document_id for document_id, grade in relevance.items() if int(grade) > 0}
    if not relevant:
        raise ValueError("relevance must contain at least one positive grade")
    return len(set(ranking[:k]) & relevant) / len(relevant)

def ndcg_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return nDCG@k with gains 2^grade-1 and log2 discount."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    grades = [int(grade) for grade in relevance.values()]
    if any(grade < 0 for grade in grades):
        raise ValueError("relevance grades must be non-negative")
    if not any(grade > 0 for grade in grades):
        raise ValueError("relevance must contain at least one positive grade")
    dcg = sum(
        (2 ** int(relevance.get(document_id, 0)) - 1.0) / math.log2(rank + 1.0)
        for rank, document_id in enumerate(ranking[:k], start=1)
    )
    ideal = sorted(grades, reverse=True)[:k]
    idcg = sum((2**grade - 1.0) / math.log2(rank + 1.0) for rank, grade in enumerate(ideal, start=1))
    return dcg / idcg

def evaluate_rankings(
    rankings: Mapping[str, Sequence[str]],
    query_records: Sequence[Mapping[str, object]],
    *,
    k: int = 5,
) -> dict[str, float]:
    """Aggregate MRR, Recall@k, and nDCG@k over one frozen query set."""

    k = _positive_int("k", k)
    if not query_records or len({r["query_id"] for r in query_records}) != len(query_records):
        raise ValueError("queries must be nonempty with unique IDs")
    if len(rankings) != len(query_records):
        raise ValueError("rankings must contain exactly one entry per query")
    reciprocal_ranks: list[float] = []
    recalls: list[float] = []
    ndcgs: list[float] = []
    for record in query_records:
        query_id = str(record["query_id"])
        if query_id not in rankings:
            raise KeyError(f"missing ranking for {query_id}")
        relevance = dict(record["relevance"])
        ranking = tuple(rankings[query_id])
        if len(ranking) != len(set(ranking)):
            raise ValueError(f"ranking for {query_id} contains duplicates")
        reciprocal_ranks.append(reciprocal_rank(ranking, relevance))
        recalls.append(recall_at_k(ranking, relevance, k))
        ndcgs.append(ndcg_at_k(ranking, relevance, k))
    return {
        "mrr": float(np.mean(reciprocal_ranks)),
        f"recall_at_{k}": float(np.mean(recalls)),
        f"ndcg_at_{k}": float(np.mean(ndcgs)),
    }
# No search or scoring calls inside evaluation: consume the stored records.
metrics = {}
for method in ("bm25","dense_lsa","rrf","reranked"):
    rankings = {q["query_id"]:[r.document_id for r in stored[q["query_id"]][method]]
                for q in queries}
    metrics[method] = evaluate_rankings(rankings,queries,k=5)
    assert all(np.isclose(value, REFERENCE["retrieval_results"][LANG]["metrics"][method][name])
               for name,value in metrics[method].items())
print(json.dumps(metrics,indent=2))
assert len(stored) == 7 and sum(len(row["reranked"]) for row in stored.values()) == 56

@dataclass(frozen=True)
class ContextItem:
    """One selected chunk and its explicit budget contribution."""

    citation: str
    document_id: str
    source_id: str
    token_count: int
    score: float
    text: str

@dataclass(frozen=True)
class ContextPack:
    """A budgeted retrieval context with selected and skipped trace data."""

    items: tuple[ContextItem, ...]
    used_tokens: int
    max_tokens: int
    skipped: Mapping[str, str]
    rendered_text: str

def pack_context(
    ranked_ids: Sequence[str],
    scores: Mapping[str, float],
    documents: Mapping[str, Mapping[str, object]],
    *,
    max_tokens: int,
    max_per_source: int = 1,
    citation_overhead_tokens: int | None = None,
) -> ContextPack:
    """Deduplicate and pack whole chunks under an explicit token budget.

    Exact normalized-text duplicates and per-source overflow are skipped.  Chunks are
    never silently truncated; a chunk either fits in full or is traced as skipped.
    """

    max_tokens = _positive_int("max_tokens", max_tokens)
    max_per_source = _positive_int("max_per_source", max_per_source)
    if citation_overhead_tokens is not None:
        raise ValueError("count the serialized citation instead of a guessed overhead")
    _validate_ids(ranked_ids)
    ranked_ids = tuple(ranked_ids)
    if not ranked_ids:
        raise ValueError("ranked_ids must be nonempty")
    selected: list[ContextItem] = []
    skipped: dict[str, str] = {}
    seen_texts: set[str] = set()
    per_source: Counter[str] = Counter()
    used = 0
    for document_id in ranked_ids:
        if document_id not in documents or document_id not in scores:
            raise KeyError(f"missing document or score for {document_id}")
        record = documents[document_id]
        text = record["text"]
        source_id = record["source_id"]
        _validate_ids([source_id])
        if not math.isfinite(float(scores[document_id])):
            raise ValueError("context scores must be finite")
        tokens = tokenize(text)
        if not tokens:
            skipped[document_id] = "empty"
            continue
        normalized = " ".join(tokens)
        if normalized in seen_texts:
            skipped[document_id] = "duplicate_text"
            continue
        if per_source[source_id] >= max_per_source:
            skipped[document_id] = "source_limit"
            continue
        citation = f"[{source_id}:{document_id}]"
        cost = len(tokenize(citation + " " + text))
        if used + cost > max_tokens:
            skipped[document_id] = "budget"
            continue
        citation = f"[{source_id}:{document_id}]"
        selected.append(
            ContextItem(
                citation=citation,
                document_id=document_id,
                source_id=source_id,
                token_count=cost,
                score=float(scores[document_id]),
                text=text,
            )
        )
        used += cost
        seen_texts.add(normalized)
        per_source[source_id] += 1
    if not selected:
        raise ValueError("no chunk fits the context budget")
    rendered = "\n\n".join(f"{item.citation} {item.text}" for item in selected)
    assert len(tokenize(rendered)) == used
    return ContextPack(tuple(selected), used, max_tokens, dict(skipped), rendered)
head = stored["q6"]["reranked"]
pack = pack_context([r.document_id for r in head],{r.document_id:r.score for r in head},
                    documents,max_tokens=42)
assert len(tokenize(pack.rendered_text)) == pack.used_tokens <= 42
expected = REFERENCE["retrieval_results"][LANG]["context_pack_q6"]
assert [r.document_id for r in pack.items] == expected["selected_ids"]
print(pack.rendered_text)
print(pack.used_tokens, pack.skipped)

all_chunks = [c for doc in data["documents"]
              for c in chunk_document(doc["document_id"],doc["text"],chunk_size=12,overlap=3)]
chunk_docs = {c.chunk_id:asdict(c) for c in all_chunks}
chunk_ids = list(chunk_docs)
chunk_texts = [chunk_docs[i]["text"] for i in chunk_ids]
lex = BM25Index(chunk_ids,chunk_texts)
vec = LSAIndex(chunk_ids,chunk_texts,dimension=4)
query = queries[6]["text"]
# Two search calls, one bounded rerank, one pack; metrics never repeat them.
branches = {"lex":[r.document_id for r in lex.search(query,top_k=8)],
            "lsa":[r.document_id for r in vec.search(query,top_k=8)]}
fused = reciprocal_rank_fusion(branches)
ranked = rerank_candidates(query,[r.document_id for r in fused[:8]],
                           dict(zip(chunk_ids,chunk_texts)))
context = pack_context([r.document_id for r in ranked],{r.document_id:r.score for r in ranked},
                       chunk_docs,max_tokens=42)
item = context.items[0]
answer = {"text":item.text, "citation":item.citation, "chunk_id":item.document_id}
record = chunk_docs[answer["chunk_id"]]
source = documents[record["source_id"]]["text"]
assert answer["text"] == source[record["start_char"]:record["end_char"]]
assert tokenize(answer["text"]) == tokenize(source)[record["start_token"]:record["end_token"]]
trace = {"search_calls":2,"rerank_calls":1,"scored_pairs":len(ranked),
         "answer":answer,"used_word_units":context.used_tokens}
assert json.loads(json.dumps(trace)) == trace
print(json.dumps(trace,ensure_ascii=False,indent=2))

def index_storage_gib(num_vectors: int, dimension: int, bits_per_component: float) -> float:
    """Return raw vector payload in GiB using powers-of-two units."""

    num_vectors = _positive_int("num_vectors", num_vectors)
    dimension = _positive_int("dimension", dimension)
    bits_per_component = _positive_finite("bits_per_component", bits_per_component)
    return num_vectors * dimension * bits_per_component / 8.0 / 2**30

def brute_ann_scenario(
    *,
    num_vectors: int,
    dimension: int,
    cpu_flops_per_second: float,
    visited_vectors: int,
) -> dict[str, float]:
    """Return a transparent exact-vs-ANN arithmetic scenario."""

    num_vectors = _positive_int("num_vectors", num_vectors)
    dimension = _positive_int("dimension", dimension)
    visited_vectors = _positive_int("visited_vectors", visited_vectors)
    cpu_flops_per_second = _positive_finite("cpu_flops_per_second", cpu_flops_per_second)
    exact_flops = 2.0 * num_vectors * dimension
    ann_flops = 2.0 * visited_vectors * dimension
    exact_seconds = exact_flops / cpu_flops_per_second
    ann_milliseconds = ann_flops / cpu_flops_per_second * 1e3
    return {
        "exact_flops": exact_flops,
        "exact_seconds": exact_seconds,
        "ann_flops": ann_flops,
        "ann_milliseconds": ann_milliseconds,
        "speedup": exact_seconds * 1e3 / ann_milliseconds,
    }

def cascade_latency_scenario() -> dict[str, float]:
    """Return a scenario mixing 2PT compute estimates and an assumed decode rate."""

    peak = 990e12
    mfu = 0.5
    effective = peak * mfu
    encoder_parameters = 0.6e9
    query_tokens = 32
    candidates = 100
    rerank_tokens = 432
    generator_parameters = 8.03e9
    context_tokens = 512 + 8 * 400
    answer_tokens = 300
    decode_tokens_per_second = 209.4
    stages = {
        "embed_ms": 2 * encoder_parameters * query_tokens / effective * 1e3,
        "ann_ms": brute_ann_scenario(
            num_vectors=10_000_000,
            dimension=1024,
            cpu_flops_per_second=10e9,
            visited_vectors=2_000,
        )["ann_milliseconds"],
        "rerank_ms": candidates * 2 * encoder_parameters * rerank_tokens / effective * 1e3,
        "prefill_ms": 2 * generator_parameters * context_tokens / effective * 1e3,
        "decode_ms": answer_tokens / decode_tokens_per_second * 1e3,
    }
    total = sum(stages.values())
    return {
        **stages,
        "total_ms": total,
        "context_tokens": context_tokens,
        "decode_share": stages["decode_ms"] / total,
        "retrieval_before_prefill_share": (
            stages["embed_ms"] + stages["ann_ms"] + stages["rerank_ms"]
        )
        / total,
    }

def long_context_vs_rag(*, long_tokens: int = 200_000, rag_tokens: int = 3_712) -> dict[str, float]:
    """Compare prefill, KV payload, and token count under one explicit model geometry."""

    long_tokens = _positive_int("long_tokens", long_tokens)
    rag_tokens = _positive_int("rag_tokens", rag_tokens)
    if rag_tokens >= long_tokens:
        raise ValueError("rag_tokens must be smaller than long_tokens")
    generator_parameters = 8.03e9
    effective_flops = 990e12 * 0.5
    kv_bytes_per_token = 128 * 2**10
    long_prefill_seconds = 2 * generator_parameters * long_tokens / effective_flops
    rag_prefill_seconds = 2 * generator_parameters * rag_tokens / effective_flops
    ratio = long_tokens / rag_tokens
    return {
        "long_tokens": long_tokens,
        "rag_tokens": rag_tokens,
        "long_prefill_seconds": long_prefill_seconds,
        "rag_prefill_seconds": rag_prefill_seconds,
        "long_kv_gib": long_tokens * kv_bytes_per_token / 2**30,
        "rag_kv_gib": rag_tokens * kv_bytes_per_token / 2**30,
        "ratio": ratio,
        "kv_bytes_per_token": kv_bytes_per_token,
    }
storage = index_storage_gib(10_000_000,1024,1)
assert storage == 1_280_000_000 / 2**30
assert analytical_chunk_count(20_000_000,400,64) == 59_524
cascade = cascade_latency_scenario()
assert np.isclose(cascade["total_ms"],1658.3129827096177)
assert np.isclose(long_context_vs_rag()["ratio"],200_000/3712)
print(json.dumps(cascade,indent=2))

def rejects(fn):
    try:
        fn()
    except (ValueError,TypeError,KeyError):
        return
    raise AssertionError("invalid input was accepted")
rejects(lambda: bm25.search("zzzzunseen"))
rejects(lambda: dense.search("zzzzunseen"))
rejects(lambda: chunk_document("x","text",chunk_size=4,overlap=4))
rejects(lambda: reciprocal_rank_fusion({"x":["a","a"]}))
rejects(lambda: ndcg_at_k(["a","a"],{"a":3},2))
rejects(lambda: ndcg_at_k(["a"],{"a":1.9},1))
rejects(lambda: evaluate_rankings({},[],k=5))
rejects(lambda: BM25Index([1,"1"],["a","b"]))
rejects(lambda: rerank_candidates("a",[1],{"1":"a"}))
rejects(lambda: pack_context([ids[0]],{ids[0]:float("nan")},documents,max_tokens=50))
rejects(lambda: pack_context([ids[0]],{ids[0]:1},documents,max_tokens=1))
assert reciprocal_rank([],{"a":1}) == 0
assert recall_at_k(["a"],{"a":1,"b":1},1) == 0.5
assert ndcg_at_k(["a","b"],{"a":3,"b":1},2) == 1
assert analytical_chunk_count(10**20+1,4,1) == (10**20+1-4+2)//3+1
print("PASS")

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"воспроизводимый тестовый запуск; seed={SEED}")
DEMO_LANGUAGE = 'ru'

import math, re, json
from collections import Counter
from dataclasses import dataclass
from typing import Mapping, Sequence
WORD_RE = re.compile(r"[^\W_]+", re.UNICODE)
def _positive_int(name: str, value: int) -> int:
    """Validate and return one strictly positive integer."""

    if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)

def _positive_finite(name: str, value: float) -> float:
    """Validate and return one strictly positive finite float."""

    value = float(value)
    if not math.isfinite(value) or value <= 0.0:
        raise ValueError(f"{name} must be finite and positive")
    return value

def tokenize(text: str) -> tuple[str, ...]:
    """Find word spans in the original text, then casefold each word for search."""

    if not isinstance(text, str):
        raise TypeError("text must be a string")
    return tuple(match.group().casefold() for match in WORD_RE.finditer(text))

def _validate_ids(ids):
    """IDs remain nonempty strings; coercion must not merge 1 and '1'."""
    if any(not isinstance(i, str) or not i.strip() for i in ids):
        raise ValueError("document IDs must be nonempty strings")
    if len(ids) != len(set(ids)):
        raise ValueError("document IDs must be unique")

def _validate_judgments(ranking, relevance):
    """Reject duplicate ranks, fractional/negative grades and empty gold sets."""
    _validate_ids(ranking)
    _validate_ids(list(relevance))
    if any(isinstance(g, bool) or not isinstance(g, (int, np.integer))
           or not 0 <= g <= 30 for g in relevance.values()):
        raise ValueError("grades must be integers from 0 through 30")
    if not any(g > 0 for g in relevance.values()):
        raise ValueError("at least one positive judgment is required")

def stable_ranking(scores: Sequence[float], document_ids: Sequence[str]) -> tuple[str, ...]:
    """Sort scores descending with document ID as a deterministic tie-break."""

    score_array = np.asarray(scores, dtype=np.float64)
    if score_array.ndim != 1 or score_array.size != len(document_ids):
        raise ValueError("scores and document_ids must be aligned rank-1 sequences")
    if not np.all(np.isfinite(score_array)):
        raise ValueError("scores contain NaN or infinity")
    _validate_ids(document_ids)
    if len(set(document_ids)) != len(document_ids):
        raise ValueError("document_ids must be unique")
    return tuple(
        document_id
        for _, document_id in sorted(
            zip(score_array.tolist(), map(str, document_ids)),
            key=lambda pair: (-pair[0], pair[1]),
        )
    )

@dataclass(frozen=True)
class SearchResult:
    """One ranked document and its branch score."""

    document_id: str
    score: float
    rank: int

class BM25Index:
    """A small exact BM25 index with explicit empty-input policy."""

    def __init__(
        self,
        document_ids: Sequence[str],
        texts: Sequence[str],
        *,
        k1: float = 1.5,
        b: float = 0.75,
    ) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        self.tokens = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in self.tokens):
            raise ValueError("empty documents must be removed before indexing")
        self.k1 = _positive_finite("k1", k1)
        self.b = float(b)
        if not math.isfinite(self.b) or not 0.0 <= self.b <= 1.0:
            raise ValueError("b must lie in [0,1]")
        n_documents = len(self.tokens)
        self.average_length = sum(map(len, self.tokens)) / n_documents
        self.document_frequency = Counter(
            token for document_tokens in self.tokens for token in set(document_tokens)
        )
        self.idf = {
            token: math.log(
                (n_documents - frequency + 0.5) / (frequency + 0.5) + 1.0
            )
            for token, frequency in self.document_frequency.items()
        }

    def scores(self, query: str) -> np.ndarray:
        """Return one BM25 score per indexed document."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        if not any(token in self.idf for token in query_tokens):
            raise ValueError("query has no indexed lexical terms")
        scores = np.zeros(len(self.tokens), dtype=np.float64)
        for row, document_tokens in enumerate(self.tokens):
            frequencies = Counter(document_tokens)
            length_norm = 1.0 - self.b + self.b * len(document_tokens) / self.average_length
            for token in query_tokens:
                tf = frequencies.get(token, 0)
                if tf == 0 or token not in self.idf:
                    continue
                scores[row] += (
                    self.idf[token]
                    * tf
                    * (self.k1 + 1.0)
                    / (tf + self.k1 * length_norm)
                )
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank all documents or the requested top-k with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )

class LSAIndex:
    """Exact cosine search over dense TF-IDF + SVD document vectors."""

    def __init__(self, document_ids: Sequence[str], texts: Sequence[str], *, dimension: int = 4) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        tokenized = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in tokenized):
            raise ValueError("empty documents must be removed before indexing")
        vocabulary = sorted({token for tokens in tokenized for token in tokens})
        self.vocabulary = {token: index for index, token in enumerate(vocabulary)}
        counts = np.zeros((len(tokenized), len(vocabulary)), dtype=np.float64)
        document_frequency = np.zeros(len(vocabulary), dtype=np.float64)
        for row, tokens in enumerate(tokenized):
            frequencies = Counter(tokens)
            for token, count in frequencies.items():
                counts[row, self.vocabulary[token]] = count
                document_frequency[self.vocabulary[token]] += 1.0
        self.idf = np.log((1.0 + len(tokenized)) / (1.0 + document_frequency)) + 1.0
        tf = np.where(counts > 0.0, 1.0 + np.log(np.maximum(counts, 1.0)), 0.0)
        tfidf = tf * self.idf
        _, singular_values, vt = np.linalg.svd(tfidf, full_matrices=False)
        numerical_rank = int(np.sum(singular_values > 1e-10))
        dimension = min(_positive_int("dimension", dimension), numerical_rank)
        self.projection = vt[:dimension].T
        dense = tfidf @ self.projection
        norms = np.linalg.norm(dense, axis=1, keepdims=True)
        if np.any(norms <= 0.0):
            raise ValueError("at least one document has a zero dense representation")
        self.document_vectors = dense / norms

    def encode_query(self, query: str) -> np.ndarray:
        """Map one non-empty in-vocabulary query to a unit dense vector."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        counts = Counter(query_tokens)
        sparse = np.zeros(len(self.vocabulary), dtype=np.float64)
        for token, count in counts.items():
            position = self.vocabulary.get(token)
            if position is not None:
                sparse[position] = (1.0 + math.log(count)) * self.idf[position]
        dense = sparse @ self.projection
        norm = float(np.linalg.norm(dense))
        if norm <= 0.0:
            raise ValueError("query contains no terms represented by the dense index")
        return dense / norm

    def scores(self, query: str) -> np.ndarray:
        """Return exact cosine scores against every document vector."""

        scores = self.document_vectors @ self.encode_query(query)
        if not np.all(np.isfinite(scores)):
            raise FloatingPointError("dense search produced NaN or infinity")
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank exact dense cosine scores with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )

@dataclass(frozen=True)
class FusedResult:
    """One RRF result with a per-branch contribution trace."""

    document_id: str
    score: float
    rank: int
    contributions: Mapping[str, float]
    source_ranks: Mapping[str, int]

def reciprocal_rank_fusion(
    rankings: Mapping[str, Sequence[str]],
    *,
    k: float = 60.0,
    weights: Mapping[str, float] | None = None,
    depth: int | None = None,
) -> tuple[FusedResult, ...]:
    """Accumulate ranks in linear time, then sort unique candidates.

    Ranks are one-based. Duplicate document IDs inside one branch are rejected because
    they would double-count evidence.  The implementation builds rank maps once and
    never calls ``list.index`` inside a loop.
    """

    k = _positive_finite("k", k)
    if not rankings:
        raise ValueError("at least one ranking is required")
    if depth is not None:
        depth = _positive_int("depth", depth)
    branch_weights = {name: 1.0 for name in rankings}
    if weights is not None:
        if set(weights) != set(rankings):
            raise ValueError("weights must contain exactly the ranking branch names")
        branch_weights = {
            name: _positive_finite(f"weight[{name}]", weights[name]) for name in rankings
        }
    scores: dict[str, float] = {}
    contributions: dict[str, dict[str, float]] = {}
    source_ranks: dict[str, dict[str, int]] = {}
    for branch_name, ranking in rankings.items():
        _validate_ids(ranking)
        sequence = tuple(ranking)
        if len(sequence) != len(set(sequence)):
            raise ValueError(f"ranking {branch_name!r} contains duplicate document IDs")
        if depth is not None:
            sequence = sequence[:depth]
        for rank, document_id in enumerate(sequence, start=1):
            contribution = branch_weights[branch_name] / (k + rank)
            scores[document_id] = scores.get(document_id, 0.0) + contribution
            contributions.setdefault(document_id, {})[branch_name] = contribution
            source_ranks.setdefault(document_id, {})[branch_name] = rank
    ordered = sorted(scores, key=lambda document_id: (-scores[document_id], document_id))
    return tuple(
        FusedResult(
            document_id=document_id,
            score=float(scores[document_id]),
            rank=rank,
            contributions=dict(sorted(contributions[document_id].items())),
            source_ranks=dict(sorted(source_ranks[document_id].items())),
        )
        for rank, document_id in enumerate(ordered, start=1)
    )

def _character_trigrams(token: str) -> frozenset[str]:
    """Return boundary-aware character trigrams for one token."""

    wrapped = f"^{token}$"
    return frozenset(wrapped[index : index + 3] for index in range(max(0, len(wrapped) - 2)))

def _token_similarity(left: str, right: str) -> float:
    """Return Jaccard similarity between two token trigram sets."""

    a = _character_trigrams(left)
    b = _character_trigrams(right)
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)

def token_interaction_score(query: str, document: str) -> float:
    """Score a query/document pair with a transparent MaxSim-style proxy.

    This is not a neural cross-encoder.  It deliberately exposes the per-query-token
    maximum interaction used by late-interaction retrieval families.
    """

    query_tokens = tokenize(query)
    document_tokens = tokenize(document)
    if not query_tokens:
        raise ValueError("query must contain at least one word token")
    if not document_tokens:
        raise ValueError("document must contain at least one word token")
    return float(
        sum(
            max(_token_similarity(query_token, document_token) for document_token in document_tokens)
            for query_token in query_tokens
        )
        / len(query_tokens)
    )

def rerank_candidates(
    query: str,
    candidate_ids: Sequence[str],
    documents: Mapping[str, str],
) -> tuple[SearchResult, ...]:
    """Rerank a candidate list with the token-interaction proxy."""

    _validate_ids(candidate_ids)
    candidate_ids = tuple(candidate_ids)
    if not candidate_ids:
        raise ValueError("candidate_ids must be non-empty")
    if len(candidate_ids) != len(set(candidate_ids)):
        raise ValueError("candidate_ids must be unique")
    missing = [document_id for document_id in candidate_ids if document_id not in documents]
    if missing:
        raise KeyError(f"missing documents: {missing}")
    scores = [token_interaction_score(query, documents[document_id]) for document_id in candidate_ids]
    ranking = stable_ranking(scores, candidate_ids)
    score_by_id = dict(zip(candidate_ids, scores))
    return tuple(
        SearchResult(document_id, float(score_by_id[document_id]), rank)
        for rank, document_id in enumerate(ranking, start=1)
    )

def reciprocal_rank(ranking: Sequence[str], relevance: Mapping[str, int]) -> float:
    """Return reciprocal rank of the first positively relevant document."""
    _validate_judgments(ranking, relevance)

    if not any(int(grade) > 0 for grade in relevance.values()):
        raise ValueError("relevance must contain at least one positive grade")
    for rank, document_id in enumerate(ranking, start=1):
        if int(relevance.get(document_id, 0)) > 0:
            return 1.0 / rank
    return 0.0

def recall_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return set recall among all positively relevant documents."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    relevant = {document_id for document_id, grade in relevance.items() if int(grade) > 0}
    if not relevant:
        raise ValueError("relevance must contain at least one positive grade")
    return len(set(ranking[:k]) & relevant) / len(relevant)

def ndcg_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return nDCG@k with gains 2^grade-1 and log2 discount."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    grades = [int(grade) for grade in relevance.values()]
    if any(grade < 0 for grade in grades):
        raise ValueError("relevance grades must be non-negative")
    if not any(grade > 0 for grade in grades):
        raise ValueError("relevance must contain at least one positive grade")
    dcg = sum(
        (2 ** int(relevance.get(document_id, 0)) - 1.0) / math.log2(rank + 1.0)
        for rank, document_id in enumerate(ranking[:k], start=1)
    )
    ideal = sorted(grades, reverse=True)[:k]
    idcg = sum((2**grade - 1.0) / math.log2(rank + 1.0) for rank, grade in enumerate(ideal, start=1))
    return dcg / idcg

def evaluate_rankings(
    rankings: Mapping[str, Sequence[str]],
    query_records: Sequence[Mapping[str, object]],
    *,
    k: int = 5,
) -> dict[str, float]:
    """Aggregate MRR, Recall@k, and nDCG@k over one frozen query set."""

    k = _positive_int("k", k)
    if not query_records or len({r["query_id"] for r in query_records}) != len(query_records):
        raise ValueError("queries must be nonempty with unique IDs")
    if len(rankings) != len(query_records):
        raise ValueError("rankings must contain exactly one entry per query")
    reciprocal_ranks: list[float] = []
    recalls: list[float] = []
    ndcgs: list[float] = []
    for record in query_records:
        query_id = str(record["query_id"])
        if query_id not in rankings:
            raise KeyError(f"missing ranking for {query_id}")
        relevance = dict(record["relevance"])
        ranking = tuple(rankings[query_id])
        if len(ranking) != len(set(ranking)):
            raise ValueError(f"ranking for {query_id} contains duplicates")
        reciprocal_ranks.append(reciprocal_rank(ranking, relevance))
        recalls.append(recall_at_k(ranking, relevance, k))
        ndcgs.append(ndcg_at_k(ranking, relevance, k))
    return {
        "mrr": float(np.mean(reciprocal_ranks)),
        f"recall_at_{k}": float(np.mean(recalls)),
        f"ndcg_at_{k}": float(np.mean(ndcgs)),
    }
texts = ["TS-999 update fails after version 4.7.3 restart",
         "TS-998 update fails after version 4.7.3 restart",
         "Restart clears the cache after an update",
         "An update changes the search index"]
ids = ["error999","error998","cache","index"]
query = "TS-999 update"
bm = BM25Index(ids,texts); ls = LSAIndex(ids,texts,dimension=2)
branches = {"bm25":[r.document_id for r in bm.search(query)],
            "lsa":[r.document_id for r in ls.search(query)]}
fused = reciprocal_rank_fusion(branches)
records = {}
for depth in (1,2,4):
    head = [r.document_id for r in fused[:depth]]
    ranked = rerank_candidates(query,head,dict(zip(ids,texts)))
    records[depth] = {"ranking":[r.document_id for r in ranked], "pairs_scored":len(ranked)}
qrels = {"error999":3}
for depth,row in records.items():
    row["recall_at_1"] = recall_at_k(row["ranking"],qrels,1)
assert sum(row["pairs_scored"] for row in records.values()) == 7
print(json.dumps({"branches":branches,"stored_results":records},indent=2))
# Work is counted as pairs, not invented milliseconds.
plt.rcParams.update({"font.family":"DejaVu Sans","font.size":12,
    "axes.titlesize":17,"axes.labelsize":13,"xtick.labelsize":11,"ytick.labelsize":11})
plt.figure(figsize=(11,5))
plt.plot(list(records),[r["recall_at_1"] for r in records.values()],"o-")
plt.xlabel("Оценённые пары" if DEMO_LANGUAGE=="ru" else "Scored pairs")
plt.ylabel("Recall@1");plt.ylim(-.05,1.1);plt.xticks([1,2,4]);plt.tight_layout();plt.show()

# An exact attention operation with designed features, not a pretrained language model.
# First row = antecedent; remaining two rows form the chunk to retrieve.
def contextualize(x):
    logits = x @ x.T / np.sqrt(x.shape[1])
    logits -= logits.max(axis=1,keepdims=True)
    weights = np.exp(logits)
    weights /= weights.sum(axis=1,keepdims=True)
    return weights @ x

river = np.array([[2.,0.],[0.,1.],[0.,1.]])
finance = np.array([[0.,2.],[0.,1.],[0.,1.]])
early = [contextualize(x[1:]).mean(axis=0) for x in (river,finance)]
late = [contextualize(x)[1:].mean(axis=0) for x in (river,finance)]
assert np.array_equal(early[0],early[1])
assert not np.allclose(late[0],late[1])
print({"early_vectors":[x.tolist() for x in early],"late_vectors":[x.tolist() for x in late]})
# Changing only the excluded antecedent changes the pooled chunk after full attention.
assert late[0][0] > late[1][0]