Tutorials · 2026-09-06 · 1,105 words · 5 min read

A minimal retrieval pipeline in Python: chunking, TF-IDF, evaluation

A complete retrieval pipeline in the standard library, run on 216 arXiv abstracts, with Recall@k and MRR measured for whole-document versus chunked indexing.

retrieval RAG evaluation python

This is the retrieval half of a retrieval-augmented pipeline with nothing hidden: a corpus loader, a chunker, an embedding function, a cosine index, and an evaluation loop, in 130 lines of standard-library Python. The embedding is TF-IDF rather than a neural model, on purpose. Every other part of the pipeline, including the evaluation, is identical whether the vectors come from term weights or from a transformer, and TF-IDF runs anywhere in a second. When you are ready to swap in a real embedding model, you change one function and keep the measurements.

The corpus is 216 arXiv abstracts pulled from the arXiv export API on 2026-09-05 and shipped as datasets/minimal-retrieval-pipeline.csv (columns: arxiv_id, published, title, abstract, url). The numbers below come from running code/minimal-retrieval-pipeline.py on that file.

What the pipeline does

Four stages, in order.

Chunking. Each abstract is split into windows of 60 words with a 20-word overlap. Abstracts shorter than 60 words become a single chunk. The chunk keeps its parent document's id, so later stages can collapse several chunks back to one document.

Embedding. Each chunk becomes a sparse vector: for each term, (1 + log tf) * idf, where idf = log((1 + N) / (1 + df)) + 1 with N the number of indexed units and df the number of units containing the term. The vector is L2-normalised so that a dot product is a cosine. That is the tf-idf weighting from Manning, Raghavan and Schütze's textbook, section 6.2.1, with smoothing so that unseen terms do not divide by zero. Stop words and single-character tokens are dropped.

Search. A query is embedded the same way. Candidate units are gathered from an inverted index (any unit sharing at least one term with the query), scored by dot product, and sorted. Chunks are collapsed to documents by keeping each document's best chunk.

Evaluation. Each paper's title is used as a query against its own abstract. A query is answered when its abstract (or a chunk of it) appears in the top k. This is a self-retrieval test: it does not need labelled queries, and it is a reasonable proxy for "can the index find the right document from a short paraphrase of it".

The core of the index is short enough to read in full:

class TfidfIndex:
    def __init__(self, docs: list[tuple[str, str]]):
        self.ids = [d for d, _ in docs]
        tfs = [Counter(tokenize(t)) for _, t in docs]
        df: Counter = Counter()
        for tf in tfs:
            df.update(tf.keys())
        n = len(docs)
        self.idf = {t: math.log((1 + n) / (1 + c)) + 1.0 for t, c in df.items()}
        self.vectors = [self._vec(tf) for tf in tfs]
        self.postings: dict[str, list[int]] = defaultdict(list)
        for i, tf in enumerate(tfs):
            for t in tf:
                self.postings[t].append(i)

    def _vec(self, tf: Counter) -> dict[str, float]:
        v = {t: (1 + math.log(c)) * self.idf.get(t, 1.0) for t, c in tf.items()}
        norm = math.sqrt(sum(x * x for x in v.values())) or 1.0
        return {t: x / norm for t, x in v.items()}

The search method collapses chunks to documents, which matters more than it looks: without the collapse, a long abstract with three chunks could take three of the top-five slots and push out a different document that should have ranked.

What it measured

Running the script on the shipped corpus prints:

corpus: 216 abstracts, 1054 chunks (60 words, 20 overlap)
whole-abstract  recall@1=0.949  recall@5=0.991  recall@10=0.991  mrr=0.969
chunked         recall@1=0.940  recall@5=0.977  recall@10=0.986  mrr=0.959

Two things to read off that table. First, title-to-abstract self-retrieval is an easy task for TF-IDF on a corpus this small; 95% of titles land their own abstract at rank one. Do not mistake that for a statement about your production corpus, where queries are not paraphrases of the target and the corpus is thousands of times larger. Second, chunking costs a little on this task, about one point of Recall@1 and 1.4 points of Recall@5. That is the expected direction: a title's terms are spread across the whole abstract, and a 60-word chunk sees only part of them, so the best chunk scores lower than the whole document would. Chunking pays off when the answer to a query sits in one paragraph of a long document and the rest of the document is noise. Abstracts are already one paragraph, so there is nothing for chunking to win here.

The script also prints a ranking for a query that is not a title, to show what the scores look like when the match is partial:

query: 'how much does long context cost compared with retrieval'
  0.165  2609.03218v1  The Analyst in the Prompt: Role, Retrieval, and Memory Biases in LLM F
  0.164  2608.28859v1  The Halt Vector: Internalizing a Causal Steering Intervention for Effi
  0.132  2609.01918v1  Grounded, Compute-Efficient LLM Policy Agents for Energy-Poverty Equit

The cosine scores are low (0.16 against a self-retrieval hit that typically scores above 0.5) and the top result is a partial lexical match on "retrieval" and "cost", not a semantic match. That is TF-IDF's known weakness and the reason to swap in a neural embedding: it matches words, not meanings. The pipeline around it does not care which.

Swapping the embedding

embed() returns a dict of term to weight, and search() computes a sparse dot product. To use a dense model, replace those two with a vector and a dense dot product, and keep the L2 normalisation so that scores remain cosines. The evaluation loop, the chunker, and the chunk-to-document collapse do not change. If you use the sentence-transformers package, batch the chunk embeddings once at index time and cache them to disk; re-embedding 1,054 chunks on every run is the kind of cost that makes people skip evaluation.

One caution when you make the swap: rerun the self-retrieval test and record the numbers before changing anything else. A dense model that scores below TF-IDF on self-retrieval is usually misconfigured (wrong pooling, missing query prefix, truncated inputs), and this test catches that in seconds.

Limitations of this test

Self-retrieval with titles is a smoke test, not a benchmark. It says whether the index is wired correctly and gives a floor for how much chunking costs. It does not measure the thing you care about in production, which is whether a user's question retrieves the passage that answers it. For that you need a small set of real queries with judged relevant passages, even 50 of them, and the same Recall@k and MRR functions in this script will score them without modification. Building that set is a separate job, and a small one: forty questions with the passage that answers each is enough to start.

The corpus is also tiny. At 216 documents the inverted index and a brute-force scan cost the same; at a million documents you need an approximate nearest-neighbour index for dense vectors and a real inverted index for sparse ones. The evaluation loop still applies, and its cost stays proportional to the number of queries, not the corpus.

The dataset

datasets/minimal-retrieval-pipeline.csv is the corpus, 216 rows, UTF-8, header row included.

_README: - arxiv_id: the arXiv identifier with version suffix, for example 2608.10484v1 - published: the submission date of that version, YYYY-MM-DD, as returned by the arXiv API - title: the paper title, whitespace-normalised - abstract: the abstract, whitespace-normalised - url: the abstract page URL

The rows were pulled through the arXiv export API's query endpoint, documented in the arXiv API Basics page, on 2026-09-05. They are recent cs.CL, cs.LG and cs.IR listings and were chosen for size and licence convenience, not for topic; the same script works on any CSV with an id and a text column.