Build retrieval over a folder of text files using TF-IDF with no external libraries. Return the top passages with scores.

“I’ll index passages, not files, because that’s what the caller gets back. Tokens are lowercased alphanumeric runs minus a few stopwords; no stemming, so plural and singular are different terms, which I’d revisit. Term frequency is sublinear, 1 + ln(tf). IDF is ln((1 + N) / (1 + df)) + 1. The outer plus-one means a term in every passage gets weight 1.0, the lowest this formula gives, instead of zero, so it can still break ties. The inner plus-ones act as if an extra passage contained every term, which damps extreme weights on a tiny corpus; they would also avoid dividing by zero for a term no passage contains, but I drop those anyway. Vectors are L2-normalized, so the dot product is cosine similarity and a long passage doesn’t win by being long.”

Sublinear term frequency and cosine normalization are as in Manning, Raghavan and Schütze, Introduction to Information Retrieval, chapter 6, whose IDF is the unsmoothed log(N / df). The smoothed form is the one scikit-learn uses with smooth_idf=True (TfidfTransformer documentation).

“Cosine also rewards short passages: a passage that is only the heading ‘Refunds’ would score perfectly against a refunds query and answer nothing. So the chunker merges short paragraphs forward into the next one.”

import heapq
import math
import re
from collections import Counter
from pathlib import Path

# Letters and digits in any script
TOKEN = re.compile(r"[^\W_]+")
STOP = {"the", "a", "an", "and", "or", "of", "to", "in",
        "is", "for", "on", "with"}


def tokens(text: str) -> list[str]:
    return [t for t in TOKEN.findall(text.lower())
            if t not in STOP]


def chunk(text: str, min_words: int = 25,
          max_words: int = 120, overlap: int = 30) -> list[str]:
    """Merge short paragraphs (headings, sign-offs) forward;
    window long ones."""
    out: list[str] = []
    buf: list[str] = []
    for para in re.split(r"\n\s*\n", text):
        buf += para.split()
        if len(buf) < min_words:
            # A bare heading waits for its paragraph
            continue
        step = max_words - overlap
        last = max(len(buf) - overlap, 1)
        for start in range(0, last, step):
            out.append(" ".join(buf[start:start + max_words]))
        buf = []
    if buf:
        # A short tail joins the last passage
        if out:
            out[-1] += " " + " ".join(buf)
        else:
            out.append(" ".join(buf))
    return out


class Index:
    def __init__(self, folder: Path):
        self.passages: list[tuple[str, int, str]] = []
        for path in sorted(folder.rglob("*.txt")):
            text = path.read_text(encoding="utf-8",
                                  errors="replace")
            # Relative path, so two faq.txt files stay distinct
            rel = str(path.relative_to(folder))
            self.passages += [(rel, i, c)
                              for i, c in enumerate(chunk(text))]

        counts = [Counter(tokens(text))
                  for _, _, text in self.passages]
        n = len(counts)
        df = Counter(term for tf in counts for term in tf)
        self.idf = {t: math.log((1 + n) / (1 + d)) + 1
                    for t, d in df.items()}

        self.vectors = [self._weigh(tf) for tf in counts]
        self.postings: dict[str, list[int]] = {}
        for i, vec in enumerate(self.vectors):
            for term in vec:
                self.postings.setdefault(term, []).append(i)

    def _weigh(self, tf: Counter) -> dict[str, float]:
        vec = {t: (1 + math.log(c)) * self.idf[t]
               for t, c in tf.items() if t in self.idf}
        norm = math.sqrt(sum(w * w for w in vec.values()))
        if not norm:
            return {}
        return {t: w / norm for t, w in vec.items()}

    def search(self, query: str, k: int = 5):
        """Top k as (score, path, chunk index, text)."""
        # Terms the corpus has never seen drop out here
        q = self._weigh(Counter(tokens(query)))
        candidates = {i for t in q for i in self.postings[t]}
        scored = (
            (sum(w * self.vectors[i].get(t, 0.0)
                 for t, w in q.items()), i)
            for i in candidates
        )
        top = heapq.nlargest(k, scored)
        return [(round(s, 4), *self.passages[i]) for s, i in top]

“Each result is score, path, chunk index and text. The path and chunk index are what a generated answer cites, so a reader can open the exact passage. Query terms the corpus has never seen have no IDF, so they drop out, and a query made only of them returns an empty list rather than raising. If the files had marked headings, I’d also prefix every window of a long section with its heading. The test I care most about is a ranking test where I know the answer:”

def test_rare_term_outranks_repeated_common_term(tmp_path):
    files = {
        "refunds.txt": "Refunds go back to the original card "
                       "after billing approves them.",
        "team.txt": "For billing, ask billing support "
                    "or the billing team.",
        "invoices.txt": "Billing sends invoices monthly.",
        "plans.txt": "Billing plans renew each year.",
    }
    for name, text in files.items():
        (tmp_path / name).write_text(text)
    top = Index(tmp_path).search("billing refunds", k=2)
    assert [path for _, path, _, _ in top] == [
        "refunds.txt", "team.txt"]
    # A wide margin, so a small tokenizer change can't flip it
    assert top[0][0] > 1.5 * top[1][0]

“‘billing’ is in every file, so IDF shrinks it, and the one file with ‘refunds’ wins even though team.txt says ‘billing’ three times. Run on that folder, the search returns:”

>>> for hit in index.search("billing refunds", k=2):
...     print(hit)
...
(0.3922, 'refunds.txt', 0, 'Refunds go back to the original card after billing approves them.')
(0.2472, 'team.txt', 0, 'For billing, ask billing support or the billing team.')
>>> index.search("chargeback")
[]

“Remove IDF and team.txt wins, 0.5454 to 0.4714, which is the regression this catches. The margin assertion is there because a test that passes by a hair can flip on a harmless stopword change, and then it no longer proves anything.

For real use I’d collect a few dozen questions from the people who’ll ask them, label the passage that answers each, and track hit rate at k every time I touch the chunker or the tokenizer. At larger scale the first thing to break is memory, because every vector sits in a Python dict; I’d move postings to disk or into a search engine with BM25, which saturates term frequency and makes length normalization tunable, and I’d reach for embeddings when queries and passages use different words for the same thing.

In a customer deployment, each passage carries its source file’s access list, and I filter candidates by the caller’s permissions before scoring, never after, so a restricted passage can’t shape the results or leak through a snippet. And when a model answers from these passages, I pass each one with its path and chunk index and require the answer to cite them, then measure retrieval (hit rate at k) and the answer’s faithfulness to its passages separately, so I know which half failed.”