bilbysearch

This repo is a work in progress - it is a package for dedicated search code in future search pipelines at Bilby.

bilbysearch turns articles into vectors - by default with intfloat/multilingual-e5-large-instruct, or with any of the 31 other models in its registry, locally or through OpenRouter - and reads and writes the daily parquet files those vectors live in.

It exists so each model's contract lives in one place. What a model expects - its output width, context window, query instruction and document prefix - is recorded once, in model_config.py, and every function takes the model by name and applies that record. Which model production uses, and pipeline rules such as the title+body join, are in production_config.py. So notebooks, batch backfills and retrieval pipelines cannot drift apart.

Install

pip install bilbysearch

That one command brings embedding and storage: the model stack (sentence-transformers, torch) and gs:// support (fsspec, gcsfs, google-cloud-storage). torch is the heavy part, at around 2 GB once CUDA wheels are involved.

Reranking with Jev is optional, because it brings bilbystats and its LLM client libraries:

pip install "bilbysearch[rerank]"

[all] installs everything. The [embed] and [gcs] extras from 0.1.x are still accepted, so an existing "bilbysearch[embed,gcs]" dependency keeps working, but they no longer add anything.

Usage

import bilbysearch

# Embed a frame of articles. Rows with no text get a null vector rather than a
# meaningless one, and three provenance columns record how the vectors were made.
articles = bilbysearch.read_parquet("daily_2026-01-31.parquet")
embedded = bilbysearch.embed_dataframe(articles)
bilbysearch.write_parquet(embedded, "daily_2026-01-31_embedding.parquet")

Queries and documents are embedded by different functions on purpose: the production model expects an instruction prefix on queries and none on documents, and skipping it measurably degrades retrieval. embed_query adds the query instruction and embed_texts the document prefix, so neither can be forgotten or put on the wrong side.

query = bilbysearch.embed_query("Chinese monetary policy easing")
passages = bilbysearch.embed_texts(["The PBoC cut the reserve requirement ratio."])
similarity = query @ passages[0]  # vectors are L2-normalised, so this is cosine

The model is loaded on first use and kept loaded - the two most recently used models stay in memory - so embedding a frame and then a query loads it once. clear_embedder_cache() releases them. Every function also takes device ("cuda", "mps" or "cpu"; the default lets sentence-transformers choose).

Choosing a model

Every function takes the model's name as model_name, defaulting to the production model. The name is all you pass: its width, context window and prefixes come from its entry in the registry.

bilbysearch.list_local_embedding_models(language="zh")  # "en", "zh" or "multilingual"
bilbysearch.get_model_spec("BAAI/bge-m3")                # what the registry records

embedded = bilbysearch.embed_dataframe(articles, model_name="intfloat/e5-large-v2")
query = bilbysearch.embed_query("PBoC easing", "intfloat/e5-large-v2")  # gets "query: "

That dataframe's articles were embedded with e5-large-v2's "passage: " prefix, and its embedding_model column says intfloat/e5-large-v2: the one name picks the model and labels the output, so the two cannot disagree. Embed queries with the same model as the documents; vectors from two models are not comparable.

Mistakes fail before anything is loaded or sent:

• An unknown name is refused, with the nearest known names suggested ("baai/bge-m3" -> did you mean 'BAAI/bge-m3'?).
• A text that already starts with its model's prefix is refused, e.g. "passage: ..." for e5-large-v2, which would otherwise be embedded as "passage: passage: ...". Pass texts without prefixes.
• A loaded model that disagrees with its registry width is refused, so its vectors never reach a vector store sized from the registry.

To switch production to another model, change MODEL_NAME in production_config.py to a local model from the registry; its width, context window and prefixes follow. tests/test_production_config.py pins the production values, so update it in the same change. Existing vectors must be re-embedded, and a vector store of a different width recreated.

To add a model, add its entry to model_config.py: its output width (measured from real output, not read from its card), context window, query instruction and document prefix (copied character for character from its card), language, whether it needs trust_remote_code, and where it runs.

Through OpenRouter

The same steps can run on OpenRouter's servers instead of this machine, for the models the registry records as hosted there (list_openrouter_embedding_models(); the production model is not one of them). Set OPENROUTER_API_KEY, then pass the model's registry name - the id OpenRouter uses and the model's prefixes come from its entry:

model_name = "intfloat/multilingual-e5-large"

embedded = bilbysearch.embed_dataframe_openrouter(articles, model_name=model_name)
passages = bilbysearch.embed_texts_openrouter(["The PBoC cut the RRR."], model_name)
query = bilbysearch.embed_query_openrouter("PBoC easing", model_name)

Locally, a text longer than the model's context window is cut short: e5-large reads only its first 512 tokens. OpenRouter would receive the whole text, so where the registry names the model's tokenizer (e5-large and the Perplexity models), each text is cut to the same window with that tokenizer before it is sent, and OpenRouter embeds exactly what the local model reads. OpenAI's and Gemini's tokenizers are not public, so their texts are sent whole.

Before mixing OpenRouter and local vectors in one index, check they match:

report = bilbysearch.compare_embeddings(local_vectors, openrouter_vectors)
report["all_match"]  # True if every pair has cosine similarity >= 0.99

check_openrouter_matches_local.py runs this check for multilingual-e5-large.

Threading

embed_texts_openrouter and embed_dataframe_openrouter send texts in batches of batch_size (default 32) per request, and use threads to keep up to max_concurrent_requests (default 8) requests in flight at once. Vectors come back in input order. Lower max_concurrent_requests if OpenRouter rate-limits you; 1 sends one request at a time.

embedded = bilbysearch.embed_dataframe_openrouter(
    articles, model_name=model_name, max_concurrent_requests=4
)

For async code, embed_texts_openrouter_async and embed_query_openrouter_async take the same arguments and can be awaited without blocking the event loop:

query = await bilbysearch.embed_query_openrouter_async("PBoC easing", model_name)

Upgrading from 0.2

0.3 takes models by name. Code written for 0.2 fails with an error saying what to pass instead, rather than running with the wrong settings:

Searching

vector_search scores every article in an embedded frame against a query and returns the closest, best first, with a cosine_similarity column. To search stored days, read them into one frame first:

uris = [bilbysearch.parquet_uri(prefix, date, "_embedding") for date in dates]
embedded = pd.concat([bilbysearch.read_parquet(uri) for uri in uris], ignore_index=True)
candidates = bilbysearch.vector_search("Chinese monetary policy easing", embedded, top_k=100)

The query is embedded with model_name - the production model unless you say otherwise - and it has to be the model that made the article vectors. Frames from embed_dataframe record that model in their embedding_model column, so a search with a different one, or over a frame that mixes days embedded by two models, is refused rather than returning a meaningless ranking.

It is brute force: every vector is loaded and compared in memory. That is exact and fine for weeks of articles; a corpus too big to load at once needs a vector store.

Reranking with Jev

Cosine similarity finds articles that are about the same things as the query. rerank_jev asks TypeSafe's Jev whether each candidate is actually relevant, and re-sorts by its calibrated probability. It needs the [rerank] extra (see Install). Pass it the top 50 to 100 candidates by cosine; the call goes through bilbystats.batch_jev_api, routed through OpenRouter by default, so it needs an OPENROUTER_API_KEY in the environment. Pass provider="typesafe" or provider="vercel" to route elsewhere, each with its own key.

candidates = bilbysearch.vector_search("Chinese monetary policy easing", embedded)
reranked = bilbysearch.rerank_jev("Chinese monetary policy easing", candidates, min_relevance=0.5)
reranked[["title", "jev_relevance", "jev_confidence", "jev_probabilities"]]

jev_relevance is the probability that the article is relevant, jev_confidence is Jev's confidence in that answer and jev_probabilities is the full yes/no dict. Because the probability is absolute, min_relevance works as a cutoff as well as a rank; ties keep the incoming cosine order.

Development

uv sync --all-extras
uv run pytest
uv run ruff check . && uv run ruff format .

The tests use stub models and a faked OpenRouter, so the suite runs in a couple of seconds, never downloads a real model and never makes an API call.

Licence

MIT — see LICENSE.