Embedding Jupyter Notebooks for Search

A notebook is a JSON document containing three kinds of content with different retrieval value, and treating it as a source file gets all three wrong. The chunking decision is different from every other file type in the repository.

What is actually in the file

An .ipynb file is JSON conforming to the nbformat schema, maintained by Project Jupyter (the nbformat format description). Version 4 is what you will encounter; the top level carries nbformat, nbformat_minor, metadata and cells.

Each cell has a cell_type of code, markdown or raw, a source field, and metadata. Code cells additionally carry execution_count and an outputs array. The one detail that catches every first implementation: source is permitted to be either a single string or a list of strings with the newlines already embedded, and both forms occur in the wild. Handle both, or half your notebooks index as the literal text of a Python list.

cell = nb["cells"][3]
src = cell["source"]
text = "".join(src) if isinstance(src, list) else src   # both are legal

Outputs are a list of objects each with an output_type, and there are four: stream (stdout or stderr, with a text field), execute_result (the value of the last expression), display_data (rich output such as a figure), and error (with ename, evalue and a traceback list). The first three carry a data object keyed by MIME type — text/plain, text/html, image/png — and the PNG value is base64.

Why cell boundaries are the wrong boundaries

The obvious scheme is one chunk per cell. It fails because a notebook distributes meaning across cells by design.

[markdown]  ## Removing outliers
            We drop rows more than 3 SD from the mean, since the
            sensor reports 9999 on a read failure.

[code]      m, s = df.value.mean(), df.value.std()
            df = df[(df.value - m).abs() <= 3 * s]

[stream]    dropped 412 rows

Chunked per cell, the code cell contains no word that appears in the query “how do we handle sensor read failures” — not “sensor”, not “failure”, not “outlier”. It is four tokens of pandas. The markdown cell above it contains every one of them but none of the code, so retrieving it gives a reader prose with no implementation. The two cells are one unit of meaning and the file separates them.

This is the mirror image of the granularity problem in ordinary source files. There, the risk is chunks that are too large and dilute the signal; here the risk is chunks that are too small to carry any. The fix is the same in shape and opposite in direction: group upwards. Attach each code cell to the markdown that precedes it, and keep grouping until the group either reaches a token budget or hits the next markdown heading.

Outputs: keep, summarise, discard

Outputs are the part unique to notebooks and the part most pipelines get wrong in one of two ways — dropping them all, which discards real signal, or keeping them all, which is much worse.

• Discard image/png and image/jpeg unconditionally. A single matplotlib figure is commonly tens of kilobytes of base64. Embedded as text it is high-entropy noise, it can be most of the notebook’s bytes, and it will dominate the token bill for the whole repository. Replace it with a marker such as [figure] and keep any caption. If figures genuinely need to be searchable that is a multimodal retrieval problem, not a text one.
• Truncate stream output. A training loop emits thousands of near-identical progress lines. The first and last few lines carry the information; the middle carries none.
• Keep error outputs, and keep them prominently. The ename and evalue pair is often the most valuable retrievable text in the entire notebook, because somebody searching ValueError: cannot reindex on an axis with duplicate labels wants exactly the cell that produced it. Keep the first and last traceback frames; drop the middle.
• Keep small text/plain results. A printed dataframe head tells you the column names, which are the domain vocabulary of the whole notebook and are frequently what somebody searches for.

Execution order is not document order

execution_count records the order in which cells were actually run, and in a working notebook it is frequently not monotonic: [1] [2] [7] [3] [12] [4] is entirely normal after a session of re-running cells. Two consequences for indexing.

First, the outputs may not correspond to the code above them. A cell showing execution_count: 3 below a cell showing 12 produced its output from a state that no longer exists. Indexing them as a coherent unit records something that was never true simultaneously. You cannot fix this, but you can detect it — check whether the sequence is monotonic and flag the notebook — and you can avoid asserting causality in the chunk text.

Second, a null execution count means the cell never ran in the saved session. Its output is absent or stale, and if a run of cells is unexecuted the notebook is a draft. That is a useful ranking signal: prefer notebooks that ran cleanly end to end, which you can approximate by the fraction of code cells with non-null counts and no error output.

A chunking scheme that works

1. Parse with the nbformat library rather than raw JSON. It normalises across minor versions and validates, so a malformed notebook fails at parse rather than producing odd chunks.
2. Walk the cells in document order, accumulating a group. Start a new group at any markdown cell whose first line is a heading, or when the current group exceeds your token budget.
3. Render each group as markdown prose, then code, then the filtered outputs, with the notebook path and the nearest enclosing heading prepended. The heading is the cheapest and most effective piece of context you can add.
4. Strip outputs per the rules above before counting tokens, or a single figure will blow the budget and split a group in a meaningless place.
5. Store cell indices in the chunk metadata so a result can link to the exact cell, and store the notebook’s content hash so the incremental index can skip it when unchanged.