Metadata-Version: 2.5
Name: bilbysearch
Version: 0.5.0
Summary: Article embeddings, vector and BM25 keyword search, and Jev reranking.
Project-URL: Homepage, https://github.com/bilbyai/bilbysearch
Project-URL: Repository, https://github.com/bilbyai/bilbysearch
Author: Samson
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.12
Requires-Dist: bm25s>=0.3.11
Requires-Dist: fsspec
Requires-Dist: gcsfs
Requires-Dist: google-cloud-storage
Requires-Dist: jieba>=0.42.1
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: pyarrow
Requires-Dist: requests
Requires-Dist: sentence-transformers>=5.4
Requires-Dist: torch
Provides-Extra: all
Requires-Dist: bilbystats[llm]>=0.2.25; extra == 'all'
Provides-Extra: embed
Provides-Extra: gcs
Provides-Extra: rerank
Requires-Dist: bilbystats[llm]>=0.2.25; extra == 'rerank'
Description-Content-Type: text/markdown

# 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](https://huggingface.co/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 searches those articles by
meaning (vector search) or by keyword (BM25, in English or Chinese), and reranks the
results with Jev.

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`](bilbysearch/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`](bilbysearch/production_config.py).
So notebooks, batch backfills and retrieval pipelines cannot drift apart.

## Install

```bash
pip install bilbysearch
```

That one command brings embedding, storage and search: the model stack
(sentence-transformers, torch), `gs://` support (fsspec, gcsfs, google-cloud-storage) and
BM25 keyword search (bm25s, and jieba to split Chinese into words). 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:

```bash
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

```python
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.

```python
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.

```python
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`](bilbysearch/production_config.py) to a local model from the
registry; its width, context window and prefixes follow.
[`tests/test_production_config.py`](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`](bilbysearch/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:

```python
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:

```python
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`](examples/embed_openrouter/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.

```python
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:

```python
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:

| 0.2 | 0.3 |
|---|---|
| `model = load_embedder(name)`, then `embed_texts(texts, model)` | `embed_texts(texts, name)` |
| `embed_query(q, model, instruction="query: ")` | `embed_query(q, name)` |
| `embed_dataframe(df, model=model, model_name=name)` | `embed_dataframe(df, model_name=name)` |
| `embed_texts(["passage: " + t for t in texts], model)` | `embed_texts(texts, name)` |
| `embed_query_openrouter(q, name, instruction="query: ")` | `embed_query_openrouter(q, name)` |
| `vector_search(q, df, model, instruction=...)` | `vector_search(q, df, name)` |
| `from bilbysearch.config import ...` | `from bilbysearch.production_config import ...`, or `from bilbysearch import ...` |
| any model id | a name in `bilbysearch.MODELS` |

## 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:

```python
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.

### Keyword search with BM25

`vector_search` matches meanings; `bm25_search` matches words. It finds exact names,
codes and figures that an embedding can blur, and misses the paraphrases an embedding
catches. It needs no embeddings, so it runs on the raw daily files. Build an index once,
then search it as often as needed:

```python
day = bilbysearch.read_parquet("daily_2026-01-31.parquet")
index = bilbysearch.build_bm25_index(day, language="zh")  # title+body, split with jieba
hits = bilbysearch.bm25_search("央行 降准", index)         # milliseconds per search
hits[["title", "bm25_score"]]
```

**Choose the language and the columns together.** A query only matches words in its own
language, and each language is split into words differently. In the daily files
`title+body` is mostly Chinese and `title_en+body_en` is English:

```python
en_index = bilbysearch.build_bm25_index(day, language="en", text_col="title_en+body_en")
hits = bilbysearch.bm25_search("reserve requirement cut", en_index)
```

- `"en"` splits text the way bm25s does by default - lowercased words of two or more
  letters or digits, English stopwords removed - so results match the team's English
  BM25 baseline.
- `"zh"` splits Chinese into words with [jieba](https://github.com/fxsjy/jieba)'s search
  mode, which also emits the shorter words inside long ones: 中国人民银行 gives 中国,
  人民, 银行 and the whole name, so a query for any of them finds it. Full-width letters
  and digits (ＧＤＰ, ４.８％) read as GDP and 4.8%, and Latin terms in Chinese text
  (PBoC, 5G) are kept.

`language` has no default, and an English index refuses a query with Chinese in it
rather than silently matching nothing.

**What comes back.** Up to `top_k` (default 100) articles that contain at least one of
the query's words, best first, with a `bm25_score` column - ready for `rerank_jev`.
Articles with none of the words are left out, so fewer than `top_k` can come back. A
score depends on the whole index as well as the article - how rare each word is across
it, how long its articles are on average - so it ranks articles within one search. It
is not comparable across indexes, and is not a relevance cutoff.

**Tuning.** `build_bm25_index` takes BM25's two parameters. `k1` (default 1.5) sets how
much repeating a word raises a score: higher lets each repeat keep adding, lower makes a
word count about the same however often it appears. `b` (default 0.75) sets how much
long articles are discounted, from 0 (not at all) to 1 (in proportion to their length).
bm25s works out each word's contribution to each article's score when the index is
built, which is what makes searches fast, so `k1` and `b` are fixed for an index: build
another to try other values.

**Cost.** Building reads every article: 2-3 s on a laptop for a day's Chinese
`title+body` (around 1,300 articles), and 0.3 s for its English translations. Each
search after that takes milliseconds. The index lives in memory.

**Known limits of `"zh"`:**

- jieba's dictionary is simplified Chinese, so traditional text is split poorly.
- A name the dictionary does not know is split into pieces - 宁德时代 gives 宁德 and 时代 -
  which match more loosely than the whole name would.
- Search mode trades some precision for recall: 人民币 (the yuan) also gives 人民, which
  is in 中国人民银行 (the People's Bank of China), so a query about the yuan also finds
  People's Bank articles, lower down.

[`ex_bm25_search.py`](examples/bm25_search/ex_bm25_search.py) runs these searches on a
few articles.

## Reranking with Jev

Cosine similarity finds articles that are *about the same things* as the query, and BM25
finds articles that *use the same words*. `rerank_jev` asks
[TypeSafe's Jev](https://developers.cloudflare.com/ai/models/typesafe/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 from
`vector_search` or `bm25_search`; 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.

```python
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 order.

## Development

```bash
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](LICENSE).
