bilbysearch 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- bilbysearch-0.1.0/.gitignore +15 -0
- bilbysearch-0.1.0/.python-version +1 -0
- bilbysearch-0.1.0/LICENSE +21 -0
- bilbysearch-0.1.0/PKG-INFO +107 -0
- bilbysearch-0.1.0/README.md +76 -0
- bilbysearch-0.1.0/bilbysearch/__init__.py +68 -0
- bilbysearch-0.1.0/bilbysearch/config.py +49 -0
- bilbysearch-0.1.0/bilbysearch/embed.py +306 -0
- bilbysearch-0.1.0/bilbysearch/py.typed +0 -0
- bilbysearch-0.1.0/bilbysearch/storage.py +258 -0
- bilbysearch-0.1.0/examples/embed/ex_embed_dataframe.py +48 -0
- bilbysearch-0.1.0/examples/embed/ex_embed_query.py +38 -0
- bilbysearch-0.1.0/examples/embed/ex_embed_texts.py +40 -0
- bilbysearch-0.1.0/examples/embed/ex_load_embedder.py +30 -0
- bilbysearch-0.1.0/examples/embed/ex_select_output_columns.py +46 -0
- bilbysearch-0.1.0/examples/embed/ex_select_text_series.py +28 -0
- bilbysearch-0.1.0/examples/storage/ex_ensure_parent_uri.py +42 -0
- bilbysearch-0.1.0/examples/storage/ex_is_gcs_uri.py +23 -0
- bilbysearch-0.1.0/examples/storage/ex_list_existing_input_dates.py +43 -0
- bilbysearch-0.1.0/examples/storage/ex_parquet_uri.py +27 -0
- bilbysearch-0.1.0/examples/storage/ex_read_parquet.py +50 -0
- bilbysearch-0.1.0/examples/storage/ex_uri_exists.py +37 -0
- bilbysearch-0.1.0/examples/storage/ex_write_parquet.py +55 -0
- bilbysearch-0.1.0/pyproject.toml +70 -0
- bilbysearch-0.1.0/tests/conftest.py +110 -0
- bilbysearch-0.1.0/tests/test_config.py +45 -0
- bilbysearch-0.1.0/tests/test_embed.py +186 -0
- bilbysearch-0.1.0/tests/test_extras.py +87 -0
- bilbysearch-0.1.0/tests/test_storage.py +167 -0
- bilbysearch-0.1.0/uv.lock +2520 -0
|
@@ -0,0 +1 @@
|
|
|
1
|
+
3.12
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Samson
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: bilbysearch
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Article embeddings with intfloat/multilingual-e5-large-instruct.
|
|
5
|
+
Project-URL: Homepage, https://github.com/bilbyai/bilbysearch
|
|
6
|
+
Project-URL: Repository, https://github.com/bilbyai/bilbysearch
|
|
7
|
+
Author: Samson
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Intended Audience :: Science/Research
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Requires-Python: >=3.12
|
|
14
|
+
Requires-Dist: numpy
|
|
15
|
+
Requires-Dist: pandas
|
|
16
|
+
Requires-Dist: pyarrow
|
|
17
|
+
Provides-Extra: all
|
|
18
|
+
Requires-Dist: fsspec; extra == 'all'
|
|
19
|
+
Requires-Dist: gcsfs; extra == 'all'
|
|
20
|
+
Requires-Dist: google-cloud-storage; extra == 'all'
|
|
21
|
+
Requires-Dist: sentence-transformers; extra == 'all'
|
|
22
|
+
Requires-Dist: torch; extra == 'all'
|
|
23
|
+
Provides-Extra: embed
|
|
24
|
+
Requires-Dist: sentence-transformers; extra == 'embed'
|
|
25
|
+
Requires-Dist: torch; extra == 'embed'
|
|
26
|
+
Provides-Extra: gcs
|
|
27
|
+
Requires-Dist: fsspec; extra == 'gcs'
|
|
28
|
+
Requires-Dist: gcsfs; extra == 'gcs'
|
|
29
|
+
Requires-Dist: google-cloud-storage; extra == 'gcs'
|
|
30
|
+
Description-Content-Type: text/markdown
|
|
31
|
+
|
|
32
|
+
# bilbysearch
|
|
33
|
+
|
|
34
|
+
**bilbysearch** turns articles into vectors with
|
|
35
|
+
[intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct),
|
|
36
|
+
and reads and writes the daily parquet files those vectors live in.
|
|
37
|
+
|
|
38
|
+
It exists so the model's contract lives in one place. The model id, the 1024-d output
|
|
39
|
+
width, the instruction prefix queries need and the `title+body` join rule are all in
|
|
40
|
+
[`config.py`](bilbysearch/config.py), so notebooks, batch backfills and retrieval
|
|
41
|
+
pipelines cannot drift apart.
|
|
42
|
+
|
|
43
|
+
## Install
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install bilbysearch[embed]
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
The base install is deliberately light — just pandas and pyarrow — so a caller that only
|
|
50
|
+
reads the parquet files, or only reads the model contract, does not pull a multi-gigabyte
|
|
51
|
+
ML stack it will never use.
|
|
52
|
+
|
|
53
|
+
| Extra | Adds | Needed for |
|
|
54
|
+
|---|---|---|
|
|
55
|
+
| *(none)* | pandas, pyarrow, numpy | reading `config`, local parquet IO |
|
|
56
|
+
| `[embed]` | sentence-transformers, torch | `load_embedder` |
|
|
57
|
+
| `[gcs]` | fsspec, gcsfs, google-cloud-storage | `gs://` URIs |
|
|
58
|
+
| `[all]` | both of the above | everything |
|
|
59
|
+
|
|
60
|
+
Until the first PyPI release, install from the repo instead:
|
|
61
|
+
|
|
62
|
+
```bash
|
|
63
|
+
uv add "bilbysearch[all] @ git+https://github.com/bilbyai/bilbysearch.git"
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
## Usage
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
import bilbysearch
|
|
70
|
+
|
|
71
|
+
model = bilbysearch.load_embedder()
|
|
72
|
+
|
|
73
|
+
# Embed a frame of articles. Rows with no text get a null vector rather than a
|
|
74
|
+
# meaningless one, and three provenance columns record how the vectors were made.
|
|
75
|
+
articles = bilbysearch.read_parquet("daily_2026-01-31.parquet")
|
|
76
|
+
embedded = bilbysearch.embed_dataframe(articles, model=model)
|
|
77
|
+
bilbysearch.write_parquet(embedded, "daily_2026-01-31_embedding.parquet")
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
Queries and passages are embedded by **different functions on purpose**: this model
|
|
81
|
+
expects an instruction prefix on queries and none on passages. Skipping it measurably
|
|
82
|
+
degrades retrieval.
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
query = bilbysearch.embed_query("Chinese monetary policy easing", model=model)
|
|
86
|
+
passages = bilbysearch.embed_texts(["The PBoC cut the reserve requirement ratio."], model=model)
|
|
87
|
+
similarity = query @ passages[0] # vectors are L2-normalised, so this is cosine
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
Only `load_embedder` needs the `[embed]` extra. Every other function works on any object
|
|
91
|
+
exposing sentence-transformers' `encode`, so a caller that already has a model can pass
|
|
92
|
+
it straight in.
|
|
93
|
+
|
|
94
|
+
## Development
|
|
95
|
+
|
|
96
|
+
```bash
|
|
97
|
+
uv sync
|
|
98
|
+
uv run pytest
|
|
99
|
+
uv run ruff check . && uv run ruff format .
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
The tests use a stub model, so the suite runs in well under a second and never downloads
|
|
103
|
+
the real 2.2 GB one.
|
|
104
|
+
|
|
105
|
+
## Licence
|
|
106
|
+
|
|
107
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
# bilbysearch
|
|
2
|
+
|
|
3
|
+
**bilbysearch** turns articles into vectors with
|
|
4
|
+
[intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct),
|
|
5
|
+
and reads and writes the daily parquet files those vectors live in.
|
|
6
|
+
|
|
7
|
+
It exists so the model's contract lives in one place. The model id, the 1024-d output
|
|
8
|
+
width, the instruction prefix queries need and the `title+body` join rule are all in
|
|
9
|
+
[`config.py`](bilbysearch/config.py), so notebooks, batch backfills and retrieval
|
|
10
|
+
pipelines cannot drift apart.
|
|
11
|
+
|
|
12
|
+
## Install
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
pip install bilbysearch[embed]
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
The base install is deliberately light — just pandas and pyarrow — so a caller that only
|
|
19
|
+
reads the parquet files, or only reads the model contract, does not pull a multi-gigabyte
|
|
20
|
+
ML stack it will never use.
|
|
21
|
+
|
|
22
|
+
| Extra | Adds | Needed for |
|
|
23
|
+
|---|---|---|
|
|
24
|
+
| *(none)* | pandas, pyarrow, numpy | reading `config`, local parquet IO |
|
|
25
|
+
| `[embed]` | sentence-transformers, torch | `load_embedder` |
|
|
26
|
+
| `[gcs]` | fsspec, gcsfs, google-cloud-storage | `gs://` URIs |
|
|
27
|
+
| `[all]` | both of the above | everything |
|
|
28
|
+
|
|
29
|
+
Until the first PyPI release, install from the repo instead:
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
uv add "bilbysearch[all] @ git+https://github.com/bilbyai/bilbysearch.git"
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## Usage
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
import bilbysearch
|
|
39
|
+
|
|
40
|
+
model = bilbysearch.load_embedder()
|
|
41
|
+
|
|
42
|
+
# Embed a frame of articles. Rows with no text get a null vector rather than a
|
|
43
|
+
# meaningless one, and three provenance columns record how the vectors were made.
|
|
44
|
+
articles = bilbysearch.read_parquet("daily_2026-01-31.parquet")
|
|
45
|
+
embedded = bilbysearch.embed_dataframe(articles, model=model)
|
|
46
|
+
bilbysearch.write_parquet(embedded, "daily_2026-01-31_embedding.parquet")
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
Queries and passages are embedded by **different functions on purpose**: this model
|
|
50
|
+
expects an instruction prefix on queries and none on passages. Skipping it measurably
|
|
51
|
+
degrades retrieval.
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
query = bilbysearch.embed_query("Chinese monetary policy easing", model=model)
|
|
55
|
+
passages = bilbysearch.embed_texts(["The PBoC cut the reserve requirement ratio."], model=model)
|
|
56
|
+
similarity = query @ passages[0] # vectors are L2-normalised, so this is cosine
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Only `load_embedder` needs the `[embed]` extra. Every other function works on any object
|
|
60
|
+
exposing sentence-transformers' `encode`, so a caller that already has a model can pass
|
|
61
|
+
it straight in.
|
|
62
|
+
|
|
63
|
+
## Development
|
|
64
|
+
|
|
65
|
+
```bash
|
|
66
|
+
uv sync
|
|
67
|
+
uv run pytest
|
|
68
|
+
uv run ruff check . && uv run ruff format .
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
The tests use a stub model, so the suite runs in well under a second and never downloads
|
|
72
|
+
the real 2.2 GB one.
|
|
73
|
+
|
|
74
|
+
## Licence
|
|
75
|
+
|
|
76
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Embedding articles with intfloat/multilingual-e5-large-instruct.
|
|
3
|
+
|
|
4
|
+
The functions callers need are re-exported here, so an importing module reads
|
|
5
|
+
|
|
6
|
+
from bilbysearch import load_embedder, embed_texts, embed_query
|
|
7
|
+
|
|
8
|
+
`config.py` holds the model's identity and input contract, so every caller -
|
|
9
|
+
notebooks, batch backfills, retrieval pipelines - embeds the same way.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from .config import (
|
|
13
|
+
COMBINED_FIELD_SEP,
|
|
14
|
+
DEFAULT_BATCH_SIZE,
|
|
15
|
+
DEFAULT_EMBEDDINGS_COL,
|
|
16
|
+
DEFAULT_NORMALIZE_EMBEDDINGS,
|
|
17
|
+
DEFAULT_TEXT_COL,
|
|
18
|
+
EMBEDDING_DIM,
|
|
19
|
+
MAX_TOKENS,
|
|
20
|
+
MODEL_NAME,
|
|
21
|
+
OUTPUT_FILENAME_SUFFIX,
|
|
22
|
+
PROVENANCE_COLUMNS,
|
|
23
|
+
QUERY_INSTRUCTION,
|
|
24
|
+
)
|
|
25
|
+
from .embed import (
|
|
26
|
+
embed_dataframe,
|
|
27
|
+
embed_query,
|
|
28
|
+
embed_texts,
|
|
29
|
+
load_embedder,
|
|
30
|
+
select_output_columns,
|
|
31
|
+
select_text_series,
|
|
32
|
+
)
|
|
33
|
+
from .storage import (
|
|
34
|
+
ensure_parent_uri,
|
|
35
|
+
is_gcs_uri,
|
|
36
|
+
list_existing_input_dates,
|
|
37
|
+
parquet_uri,
|
|
38
|
+
read_parquet,
|
|
39
|
+
uri_exists,
|
|
40
|
+
write_parquet,
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
__all__ = [
|
|
44
|
+
"COMBINED_FIELD_SEP",
|
|
45
|
+
"DEFAULT_BATCH_SIZE",
|
|
46
|
+
"DEFAULT_EMBEDDINGS_COL",
|
|
47
|
+
"DEFAULT_NORMALIZE_EMBEDDINGS",
|
|
48
|
+
"DEFAULT_TEXT_COL",
|
|
49
|
+
"EMBEDDING_DIM",
|
|
50
|
+
"MAX_TOKENS",
|
|
51
|
+
"MODEL_NAME",
|
|
52
|
+
"OUTPUT_FILENAME_SUFFIX",
|
|
53
|
+
"PROVENANCE_COLUMNS",
|
|
54
|
+
"QUERY_INSTRUCTION",
|
|
55
|
+
"embed_dataframe",
|
|
56
|
+
"embed_query",
|
|
57
|
+
"embed_texts",
|
|
58
|
+
"ensure_parent_uri",
|
|
59
|
+
"is_gcs_uri",
|
|
60
|
+
"list_existing_input_dates",
|
|
61
|
+
"load_embedder",
|
|
62
|
+
"parquet_uri",
|
|
63
|
+
"read_parquet",
|
|
64
|
+
"select_output_columns",
|
|
65
|
+
"select_text_series",
|
|
66
|
+
"uri_exists",
|
|
67
|
+
"write_parquet",
|
|
68
|
+
]
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Settings for the multilingual-e5-large-instruct embedding model.
|
|
3
|
+
|
|
4
|
+
These values are the single source of truth for the model's identity and its
|
|
5
|
+
input contract. Every caller - `embed.py`, notebooks, batch backfills,
|
|
6
|
+
retrieval pipelines - imports them from here, so no two callers can drift apart.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
# HuggingFace model id. Downloaded on first use and cached in ~/.cache/huggingface.
|
|
10
|
+
MODEL_NAME = "intfloat/multilingual-e5-large-instruct"
|
|
11
|
+
|
|
12
|
+
# Output vector width. The Qdrant collection is created with this size, so a model
|
|
13
|
+
# swap that changes it requires recreating the collection before ingesting.
|
|
14
|
+
EMBEDDING_DIM = 1024
|
|
15
|
+
|
|
16
|
+
# The model's context window. sentence-transformers truncates silently at this many
|
|
17
|
+
# tokens; we do not truncate ourselves.
|
|
18
|
+
MAX_TOKENS = 512
|
|
19
|
+
|
|
20
|
+
# Queries take an instruction; passages do not. Omitting this on the query side
|
|
21
|
+
# degrades retrieval quality significantly, so `embed_query` applies it and
|
|
22
|
+
# `embed_texts` does not.
|
|
23
|
+
QUERY_INSTRUCTION = (
|
|
24
|
+
"Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: "
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
# Inserted between fields when a text column is "+"-separated, e.g. "title+body".
|
|
28
|
+
COMBINED_FIELD_SEP = "\n\n"
|
|
29
|
+
|
|
30
|
+
# Columns embedded per article in production.
|
|
31
|
+
DEFAULT_TEXT_COL = "title+body"
|
|
32
|
+
|
|
33
|
+
# L2-normalize vectors by default, for cosine/dot-product search. A near no-op for
|
|
34
|
+
# multilingual-e5-large-instruct, whose pipeline already ends in a Normalize module.
|
|
35
|
+
DEFAULT_NORMALIZE_EMBEDDINGS = True
|
|
36
|
+
|
|
37
|
+
# Column embed_dataframe writes the vectors to.
|
|
38
|
+
DEFAULT_EMBEDDINGS_COL = "embeddings"
|
|
39
|
+
|
|
40
|
+
# Columns embed_dataframe writes beside the vectors, recording how they were made.
|
|
41
|
+
# Named here so the backfill can select its output schema without restating them.
|
|
42
|
+
PROVENANCE_COLUMNS = ("embedding_model", "embedding_text_col", "embedding_normalized")
|
|
43
|
+
|
|
44
|
+
# Appended to the stem of backfill output files, so an embedded day is
|
|
45
|
+
# daily_YYYY-MM-DD_embedding.parquet next to the daily_YYYY-MM-DD.parquet it came from.
|
|
46
|
+
OUTPUT_FILENAME_SUFFIX = "_embedding"
|
|
47
|
+
|
|
48
|
+
# Modest default suited to a laptop. The backfill passes a much larger value.
|
|
49
|
+
DEFAULT_BATCH_SIZE = 32
|
|
@@ -0,0 +1,306 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Loading the embedding model and turning text into vectors.
|
|
3
|
+
|
|
4
|
+
Documents and queries are embedded by different functions on purpose:
|
|
5
|
+
multilingual-e5-large-instruct expects an instruction prefix on queries and no
|
|
6
|
+
prefix on passages (see `config.QUERY_INSTRUCTION`).
|
|
7
|
+
|
|
8
|
+
Only `load_embedder` needs sentence-transformers installed (`pip install
|
|
9
|
+
bilbysearch[embed]`). The rest of this module works on any object exposing
|
|
10
|
+
SentenceTransformer's `encode`, so a caller that already has a model can pass it in.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import logging
|
|
16
|
+
|
|
17
|
+
import numpy as np
|
|
18
|
+
import pandas as pd
|
|
19
|
+
|
|
20
|
+
try: # ships with the optional [embed] extra
|
|
21
|
+
from sentence_transformers import SentenceTransformer
|
|
22
|
+
except ImportError: # pragma: no cover - exercised by the blocked-import test
|
|
23
|
+
SentenceTransformer = None
|
|
24
|
+
|
|
25
|
+
from .config import (
|
|
26
|
+
COMBINED_FIELD_SEP,
|
|
27
|
+
DEFAULT_BATCH_SIZE,
|
|
28
|
+
DEFAULT_EMBEDDINGS_COL,
|
|
29
|
+
DEFAULT_NORMALIZE_EMBEDDINGS,
|
|
30
|
+
DEFAULT_TEXT_COL,
|
|
31
|
+
EMBEDDING_DIM,
|
|
32
|
+
MODEL_NAME,
|
|
33
|
+
PROVENANCE_COLUMNS,
|
|
34
|
+
QUERY_INSTRUCTION,
|
|
35
|
+
)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def load_embedder(model_name: str = MODEL_NAME, device: str | None = None) -> SentenceTransformer:
|
|
39
|
+
"""
|
|
40
|
+
Load the sentence-transformers embedding model.
|
|
41
|
+
|
|
42
|
+
Warns if the loaded model's output width differs from the expected 1024
|
|
43
|
+
dimensions, because a mismatch is rejected downstream by the vector store
|
|
44
|
+
rather than being caught here.
|
|
45
|
+
---------------------------------------------------------------------------
|
|
46
|
+
ARGUMENTS:
|
|
47
|
+
model_name : str, optional (default=config.MODEL_NAME)
|
|
48
|
+
HuggingFace model id, downloaded on first use and cached locally.
|
|
49
|
+
device : str or None, optional (default=None)
|
|
50
|
+
'cuda', 'mps' or 'cpu'. None (or 'auto') lets sentence-transformers choose.
|
|
51
|
+
---------------------------------------------------------------------------
|
|
52
|
+
OUTPUT:
|
|
53
|
+
model : SentenceTransformer
|
|
54
|
+
The loaded model, ready to pass to the embedding functions below.
|
|
55
|
+
---------------------------------------------------------------------------
|
|
56
|
+
AUTHORS: Samson
|
|
57
|
+
---------------------------------------------------------------------------
|
|
58
|
+
"""
|
|
59
|
+
if SentenceTransformer is None:
|
|
60
|
+
raise ImportError(
|
|
61
|
+
"load_embedder needs sentence-transformers, which ships with the optional "
|
|
62
|
+
"[embed] extra: pip install bilbysearch[embed]"
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
resolved_device = None if device in (None, "auto") else device
|
|
66
|
+
logging.info("Loading %s on %s", model_name, resolved_device or "auto")
|
|
67
|
+
model = SentenceTransformer(model_name, device=resolved_device)
|
|
68
|
+
|
|
69
|
+
dim = model.get_embedding_dimension()
|
|
70
|
+
if dim != EMBEDDING_DIM:
|
|
71
|
+
logging.warning(
|
|
72
|
+
"%s produces %s-d vectors, but %s-d was expected; a vector store created "
|
|
73
|
+
"for the expected width will reject these.",
|
|
74
|
+
model_name,
|
|
75
|
+
dim,
|
|
76
|
+
EMBEDDING_DIM,
|
|
77
|
+
)
|
|
78
|
+
return model
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def embed_texts(
|
|
82
|
+
texts: list[str],
|
|
83
|
+
model: SentenceTransformer,
|
|
84
|
+
batch_size: int = DEFAULT_BATCH_SIZE,
|
|
85
|
+
normalize_embeddings: bool = DEFAULT_NORMALIZE_EMBEDDINGS,
|
|
86
|
+
show_progress_bar: bool = False,
|
|
87
|
+
) -> np.ndarray:
|
|
88
|
+
"""
|
|
89
|
+
Embed a list of documents (no instruction prefix).
|
|
90
|
+
|
|
91
|
+
Text longer than the model's context window is truncated silently by
|
|
92
|
+
sentence-transformers at `config.MAX_TOKENS` tokens.
|
|
93
|
+
---------------------------------------------------------------------------
|
|
94
|
+
ARGUMENTS:
|
|
95
|
+
texts : list of str
|
|
96
|
+
The documents to embed. May be empty.
|
|
97
|
+
model : SentenceTransformer
|
|
98
|
+
A model from `load_embedder`.
|
|
99
|
+
batch_size : int, optional (default=config.DEFAULT_BATCH_SIZE)
|
|
100
|
+
Number of texts encoded per forward pass.
|
|
101
|
+
normalize_embeddings : bool, optional (default=config.DEFAULT_NORMALIZE_EMBEDDINGS)
|
|
102
|
+
L2-normalize the output. A no-op for multilingual-e5-large-instruct,
|
|
103
|
+
whose pipeline already ends in a Normalize module.
|
|
104
|
+
show_progress_bar : bool, optional (default=False)
|
|
105
|
+
Display a progress bar while encoding.
|
|
106
|
+
---------------------------------------------------------------------------
|
|
107
|
+
OUTPUT:
|
|
108
|
+
embeddings : np.ndarray
|
|
109
|
+
Array of shape (len(texts), EMBEDDING_DIM).
|
|
110
|
+
---------------------------------------------------------------------------
|
|
111
|
+
AUTHORS: Samson
|
|
112
|
+
---------------------------------------------------------------------------
|
|
113
|
+
"""
|
|
114
|
+
if not texts:
|
|
115
|
+
return np.empty((0, EMBEDDING_DIM), dtype="float32")
|
|
116
|
+
|
|
117
|
+
return model.encode(
|
|
118
|
+
texts,
|
|
119
|
+
batch_size=batch_size,
|
|
120
|
+
show_progress_bar=show_progress_bar,
|
|
121
|
+
convert_to_numpy=True,
|
|
122
|
+
normalize_embeddings=normalize_embeddings,
|
|
123
|
+
)
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def embed_query(
|
|
127
|
+
query: str,
|
|
128
|
+
model: SentenceTransformer,
|
|
129
|
+
instruction: str = QUERY_INSTRUCTION,
|
|
130
|
+
normalize_embeddings: bool = DEFAULT_NORMALIZE_EMBEDDINGS,
|
|
131
|
+
) -> np.ndarray:
|
|
132
|
+
"""
|
|
133
|
+
Embed a search query, applying the model's instruction prefix.
|
|
134
|
+
|
|
135
|
+
Use this for the query side of a search and `embed_texts` for the articles.
|
|
136
|
+
Embedding a query without the instruction measurably degrades retrieval.
|
|
137
|
+
---------------------------------------------------------------------------
|
|
138
|
+
ARGUMENTS:
|
|
139
|
+
query : str
|
|
140
|
+
The search query, without any prefix.
|
|
141
|
+
model : SentenceTransformer
|
|
142
|
+
A model from `load_embedder`.
|
|
143
|
+
instruction : str, optional (default=config.QUERY_INSTRUCTION)
|
|
144
|
+
Prefix prepended to the query before encoding.
|
|
145
|
+
normalize_embeddings : bool, optional (default=config.DEFAULT_NORMALIZE_EMBEDDINGS)
|
|
146
|
+
L2-normalize the output (a no-op for this model).
|
|
147
|
+
---------------------------------------------------------------------------
|
|
148
|
+
OUTPUT:
|
|
149
|
+
embedding : np.ndarray
|
|
150
|
+
A single vector of shape (EMBEDDING_DIM,).
|
|
151
|
+
---------------------------------------------------------------------------
|
|
152
|
+
AUTHORS: Samson
|
|
153
|
+
---------------------------------------------------------------------------
|
|
154
|
+
"""
|
|
155
|
+
embeddings = embed_texts(
|
|
156
|
+
[instruction + query],
|
|
157
|
+
model,
|
|
158
|
+
batch_size=1,
|
|
159
|
+
normalize_embeddings=normalize_embeddings,
|
|
160
|
+
)
|
|
161
|
+
return embeddings[0]
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def select_text_series(df: pd.DataFrame, text_col: str = DEFAULT_TEXT_COL) -> pd.Series:
|
|
165
|
+
"""
|
|
166
|
+
Return the text to embed, joining columns when text_col is "+"-separated.
|
|
167
|
+
|
|
168
|
+
"title+body" concatenates those two columns with a blank line between them.
|
|
169
|
+
Missing values become empty strings rather than the literal "nan".
|
|
170
|
+
---------------------------------------------------------------------------
|
|
171
|
+
ARGUMENTS:
|
|
172
|
+
df : pd.DataFrame
|
|
173
|
+
Frame containing the named column(s).
|
|
174
|
+
text_col : str, optional (default=config.DEFAULT_TEXT_COL)
|
|
175
|
+
One column name, or several joined with "+".
|
|
176
|
+
---------------------------------------------------------------------------
|
|
177
|
+
OUTPUT:
|
|
178
|
+
texts : pd.Series
|
|
179
|
+
One string per row, aligned with df's index.
|
|
180
|
+
---------------------------------------------------------------------------
|
|
181
|
+
AUTHORS: Samson
|
|
182
|
+
---------------------------------------------------------------------------
|
|
183
|
+
"""
|
|
184
|
+
columns = text_col.split("+")
|
|
185
|
+
missing = [col for col in columns if col not in df.columns]
|
|
186
|
+
if missing:
|
|
187
|
+
raise ValueError(f"Input parquet is missing required text column(s): {missing}")
|
|
188
|
+
|
|
189
|
+
texts = df[columns[0]].fillna("").astype(str)
|
|
190
|
+
for col in columns[1:]:
|
|
191
|
+
texts = texts + COMBINED_FIELD_SEP + df[col].fillna("").astype(str)
|
|
192
|
+
return texts
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def embed_dataframe(
|
|
196
|
+
df: pd.DataFrame,
|
|
197
|
+
*,
|
|
198
|
+
model: SentenceTransformer,
|
|
199
|
+
model_name: str = MODEL_NAME,
|
|
200
|
+
text_col: str = DEFAULT_TEXT_COL,
|
|
201
|
+
embeddings_col: str = DEFAULT_EMBEDDINGS_COL,
|
|
202
|
+
batch_size: int = DEFAULT_BATCH_SIZE,
|
|
203
|
+
normalize_embeddings: bool = DEFAULT_NORMALIZE_EMBEDDINGS,
|
|
204
|
+
show_progress_bar: bool = True,
|
|
205
|
+
) -> pd.DataFrame:
|
|
206
|
+
"""
|
|
207
|
+
Add an embedding column to a frame of articles.
|
|
208
|
+
|
|
209
|
+
Rows whose combined text is empty are skipped and receive a null embedding
|
|
210
|
+
rather than a vector. Three provenance columns record how the vectors were
|
|
211
|
+
made, so an output file describes itself.
|
|
212
|
+
---------------------------------------------------------------------------
|
|
213
|
+
ARGUMENTS:
|
|
214
|
+
df : pd.DataFrame
|
|
215
|
+
Articles to embed. Not modified; a copy is returned.
|
|
216
|
+
model : SentenceTransformer
|
|
217
|
+
A model from `load_embedder`.
|
|
218
|
+
model_name : str, optional (default=config.MODEL_NAME)
|
|
219
|
+
Recorded in the embedding_model column.
|
|
220
|
+
text_col : str, optional (default=config.DEFAULT_TEXT_COL)
|
|
221
|
+
Column(s) to embed, "+"-joined. See `select_text_series`.
|
|
222
|
+
embeddings_col : str, optional (default=config.DEFAULT_EMBEDDINGS_COL)
|
|
223
|
+
Name of the column the vectors are written to.
|
|
224
|
+
batch_size : int, optional (default=config.DEFAULT_BATCH_SIZE)
|
|
225
|
+
Number of texts encoded per forward pass.
|
|
226
|
+
normalize_embeddings : bool, optional (default=config.DEFAULT_NORMALIZE_EMBEDDINGS)
|
|
227
|
+
L2-normalize the output (a no-op for this model).
|
|
228
|
+
show_progress_bar : bool, optional (default=True)
|
|
229
|
+
Display a progress bar while encoding.
|
|
230
|
+
---------------------------------------------------------------------------
|
|
231
|
+
OUTPUT:
|
|
232
|
+
out : pd.DataFrame
|
|
233
|
+
Copy of df with the embedding column plus embedding_model,
|
|
234
|
+
embedding_text_col and embedding_normalized.
|
|
235
|
+
---------------------------------------------------------------------------
|
|
236
|
+
AUTHORS: Samson
|
|
237
|
+
---------------------------------------------------------------------------
|
|
238
|
+
"""
|
|
239
|
+
out = df.copy()
|
|
240
|
+
selected = select_text_series(out, text_col)
|
|
241
|
+
valid_mask = selected.str.strip().ne("")
|
|
242
|
+
texts = selected[valid_mask].tolist()
|
|
243
|
+
|
|
244
|
+
logging.info(
|
|
245
|
+
"Embedding %s of %s rows from %r (%s skipped: empty text)",
|
|
246
|
+
len(texts),
|
|
247
|
+
len(out),
|
|
248
|
+
text_col,
|
|
249
|
+
len(out) - len(texts),
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
embeddings: list[list[float] | None] = [None] * len(out)
|
|
253
|
+
if texts:
|
|
254
|
+
encoded = embed_texts(
|
|
255
|
+
texts,
|
|
256
|
+
model,
|
|
257
|
+
batch_size=batch_size,
|
|
258
|
+
normalize_embeddings=normalize_embeddings,
|
|
259
|
+
show_progress_bar=show_progress_bar,
|
|
260
|
+
)
|
|
261
|
+
valid_positions = [idx for idx, is_valid in enumerate(valid_mask.tolist()) if is_valid]
|
|
262
|
+
for positional_index, vector in zip(valid_positions, encoded, strict=True):
|
|
263
|
+
embeddings[positional_index] = vector.astype("float32").tolist()
|
|
264
|
+
|
|
265
|
+
model_column, text_column, normalized_column = PROVENANCE_COLUMNS
|
|
266
|
+
out[embeddings_col] = embeddings
|
|
267
|
+
out[model_column] = model_name
|
|
268
|
+
out[text_column] = text_col
|
|
269
|
+
out[normalized_column] = normalize_embeddings
|
|
270
|
+
return out
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
def select_output_columns(
|
|
274
|
+
df: pd.DataFrame,
|
|
275
|
+
*,
|
|
276
|
+
id_col: str,
|
|
277
|
+
embeddings_col: str = DEFAULT_EMBEDDINGS_COL,
|
|
278
|
+
) -> pd.DataFrame:
|
|
279
|
+
"""
|
|
280
|
+
Reduce an embedded frame to the columns the backfill stores.
|
|
281
|
+
|
|
282
|
+
The stored schema is the article id plus the four columns embed_dataframe adds:
|
|
283
|
+
the vectors and the three provenance columns. Everything else in the daily file
|
|
284
|
+
stays in the daily file, which remains the source of truth for article content.
|
|
285
|
+
---------------------------------------------------------------------------
|
|
286
|
+
ARGUMENTS:
|
|
287
|
+
df : pd.DataFrame
|
|
288
|
+
A frame returned by `embed_dataframe`.
|
|
289
|
+
id_col : str
|
|
290
|
+
Column identifying the article, e.g. 'uuid'. Joins the vectors back to the
|
|
291
|
+
daily file they came from.
|
|
292
|
+
embeddings_col : str, optional (default=config.DEFAULT_EMBEDDINGS_COL)
|
|
293
|
+
Column holding the vectors.
|
|
294
|
+
---------------------------------------------------------------------------
|
|
295
|
+
OUTPUT:
|
|
296
|
+
out : pd.DataFrame
|
|
297
|
+
Copy of df with only those five columns, in that order.
|
|
298
|
+
---------------------------------------------------------------------------
|
|
299
|
+
AUTHORS: Samson
|
|
300
|
+
---------------------------------------------------------------------------
|
|
301
|
+
"""
|
|
302
|
+
columns = [id_col, embeddings_col, *PROVENANCE_COLUMNS]
|
|
303
|
+
missing = [col for col in columns if col not in df.columns]
|
|
304
|
+
if missing:
|
|
305
|
+
raise ValueError(f"Embedded frame is missing expected column(s): {missing}")
|
|
306
|
+
return df[columns].copy()
|
|
File without changes
|