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.
Files changed (30) hide show
  1. bilbysearch-0.1.0/.gitignore +15 -0
  2. bilbysearch-0.1.0/.python-version +1 -0
  3. bilbysearch-0.1.0/LICENSE +21 -0
  4. bilbysearch-0.1.0/PKG-INFO +107 -0
  5. bilbysearch-0.1.0/README.md +76 -0
  6. bilbysearch-0.1.0/bilbysearch/__init__.py +68 -0
  7. bilbysearch-0.1.0/bilbysearch/config.py +49 -0
  8. bilbysearch-0.1.0/bilbysearch/embed.py +306 -0
  9. bilbysearch-0.1.0/bilbysearch/py.typed +0 -0
  10. bilbysearch-0.1.0/bilbysearch/storage.py +258 -0
  11. bilbysearch-0.1.0/examples/embed/ex_embed_dataframe.py +48 -0
  12. bilbysearch-0.1.0/examples/embed/ex_embed_query.py +38 -0
  13. bilbysearch-0.1.0/examples/embed/ex_embed_texts.py +40 -0
  14. bilbysearch-0.1.0/examples/embed/ex_load_embedder.py +30 -0
  15. bilbysearch-0.1.0/examples/embed/ex_select_output_columns.py +46 -0
  16. bilbysearch-0.1.0/examples/embed/ex_select_text_series.py +28 -0
  17. bilbysearch-0.1.0/examples/storage/ex_ensure_parent_uri.py +42 -0
  18. bilbysearch-0.1.0/examples/storage/ex_is_gcs_uri.py +23 -0
  19. bilbysearch-0.1.0/examples/storage/ex_list_existing_input_dates.py +43 -0
  20. bilbysearch-0.1.0/examples/storage/ex_parquet_uri.py +27 -0
  21. bilbysearch-0.1.0/examples/storage/ex_read_parquet.py +50 -0
  22. bilbysearch-0.1.0/examples/storage/ex_uri_exists.py +37 -0
  23. bilbysearch-0.1.0/examples/storage/ex_write_parquet.py +55 -0
  24. bilbysearch-0.1.0/pyproject.toml +70 -0
  25. bilbysearch-0.1.0/tests/conftest.py +110 -0
  26. bilbysearch-0.1.0/tests/test_config.py +45 -0
  27. bilbysearch-0.1.0/tests/test_embed.py +186 -0
  28. bilbysearch-0.1.0/tests/test_extras.py +87 -0
  29. bilbysearch-0.1.0/tests/test_storage.py +167 -0
  30. bilbysearch-0.1.0/uv.lock +2520 -0
@@ -0,0 +1,15 @@
1
+ # Python / uv
2
+ .venv/
3
+ __pycache__/
4
+ *.py[cod]
5
+ .ruff_cache/
6
+
7
+ # Build artifacts
8
+ dist/
9
+ *.egg-info/
10
+
11
+ # OS
12
+ .DS_Store
13
+
14
+ # Claude files
15
+ CLAUDE.md
@@ -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