pymongopredict.embeddings_api

The embeddings API: datasets and a featurizer that turn documents into X for scikit-learn.

This is the scikit-learn workflow underneath $predict (see pymongopredict.predictions_api): use it to pick your own estimator, cross-validate, combine features or forecast.

What you use:

MongoDataset: the documents of a collection (with a filter, limit and target) or a list of documents, with the target values y. It indexes like an array, so scikit-learn's train_test_split, cross-validation splitters and cross_val_score work on it directly.
DocumentFeaturizer: an sklearn transformer from documents to chosen columns and embeddings of the documents, for any sklearn estimator or Pipeline.
MongoSeries: a time series aggregated from documents, for sktime forecasters.
SentenceTransformerEmbedder, VoyageEmbedder, HashingEmbedder: embedding models with non-default settings, to pass as embedding_model (a name such as "voyage-4-nano" or "voyage-4-lite" is enough otherwise).

Datasets and splits

MongoDataset(coll, target=...) reads a collection (filter=, projection=, limit= and sort= are passed to find). dataset.y is a float array for numeric targets, bool for booleans and strings otherwise; documents without a target value are left out. Split it with scikit-learn: a dataset indexes like an array (ds[indices], ds.iloc[...], ds.take(...)), so the splitters return train and test datasets that share the same source:

from sklearn.model_selection import GroupShuffleSplit, TimeSeriesSplit, train_test_split

ds_train, ds_test = train_test_split(ds, test_size=0.2, stratify=ds.y, random_state=0)  # by the target
ds_train, ds_test = train_test_split(ds, test_size=0.2, stratify=ds.field_values("address.market"))
tr, te = next(GroupShuffleSplit(test_size=0.3).split(ds, groups=ds.field_values("host.host_id")))
ds_train, ds_test = ds[tr], ds[te]                                            # a host on one side
ds = MongoDataset(coll, target="y", sort=[("ts", 1)])                         # in time order, then
ds_train, ds_test = train_test_split(ds, test_size=0.2, shuffle=False)        # train on the past
for tr, te in TimeSeriesSplit(n_splits=5, gap=24).split(ds): ...              # rolling, with a gap

cross_val_score(make_pipeline(DocumentFeaturizer(...), Ridge()), ds, ds.y) cross-validates in the same way.

Featurizing and fitting

DocumentFeaturizer.fit / transform accept a MongoDataset, a pymongo cursor, a generator, a list of documents or a DataFrame of records, and return a float32 matrix. When fitting on a MongoDataset, its target field is removed from the embedded JSON (pass exclude= for plain documents), so the answer cannot leak into the features. The featurizer is an ordinary scikit-learn transformer, so it works with Pipeline, clone and get_feature_names_out:

from sklearn.linear_model import RidgeClassifierCV
from sklearn.pipeline import make_pipeline

pipe = make_pipeline(prep, RidgeClassifierCV())
pipe.fit(ds_train, ds_train.y)
pipe.score(ds_test, ds_test.y)
pipe.predict(coll.find({"address.market": "Porto"}).limit(5))   # predict on any cursor

For cross-validation or model selection, compute X = prep.fit_transform(ds) once and pass X and ds.y to scikit-learn. Embeddings are cached, so refitting with another estimator does not embed a document again.

Forecasting with sktime

MongoSeries aggregates a time series from a collection, and series.split cuts it in time order (sktime's temporal_train_test_split(series.series, ...) does the same). The parts are pd.Series with a PeriodIndex for any sktime forecaster:

from sktime.forecasting.exp_smoothing import ExponentialSmoothing
from pymongopredict import MongoSeries

series = MongoSeries(coll_weather, time_field="ts", value_field="airTemperature.value",
                     freq="h", agg="mean", valid=lambda v: v < 900)
y_train, y_test = series.split(test_size=24)
fc = ExponentialSmoothing(trend="add", seasonal="add", sp=24).fit(y_train)
fc.predict(fh=range(1, 25))

Streaming large collections

By default a MongoDataset fetches its documents once and keeps them in memory. When documents are much larger than the features taken from them (for example listings that carry all their reviews), pass streaming=True:

ds = MongoDataset(coll, target="room_type", streaming=True, batch_size=512)

Loading reads only _id and the target, and field_values (for a stratify or group field) fetches only that field.
fit / transform fetch documents in batches of _id values, with only the fields the featurizer uses when embed_fields is given, and keep only the feature vectors.
The next batch is fetched in a background thread while the current one is embedded (prefetch=2).
Cursors and generators passed to fit / transform are also read batch by batch.

The results are the same as in memory (the tests check this). On the Airbnb demo, peak memory drops from about 300 MB to about 215 MB, and the difference grows with the size of the documents.

How documents become features

MongoDataset(collection, target, filter)  ->  sklearn train_test_split  ->  ds_train / ds_test
                                                                               |
                     <a href="#DocumentFeaturizer.fit_transform">DocumentFeaturizer.fit_transform</a>(ds_train) / .transform(ds_test)
                                                                               |
    columns (dotted paths): types inferred                 embed_fields (or the whole document):
    (number, bool, date, category, list size),             compact JSON -> embedding model
    missing values filled, scaled, one-hot encoded         (voyage-4-nano, Hugging Face, Voyage API)
                                                                               |
                                    batch by batch, side by side  ->  X for any scikit-learn estimator

Embeddings are cached in a sqlite file (.mongopredict_cache/), keyed by the model and the text.
Documents are serialised with orjson. max_json_chars (8000 by default) is applied field by field, so a document with a very long trailing array is not serialised in full only to be cut. max_tokens=32768 or max_tokens="model" measures the budget in tokens with the model's tokenizer instead.
BSON types are converted: Decimal128 to float, ObjectId to string, dates to ISO strings.

Embedding models

embedding_model= ("embeddingModel" in a prediction model definition) accepts:

value	embedding model
`"voyage-4-nano"` (default)	runs locally with sentence-transformers, 346M parameters, no API key
`"org/name"` or a path containing `/`	any sentence-transformers model from Hugging Face or disk
other `"voyage-*"` names: `"voyage-4-lite"`, `"voyage-4"`, `"voyage-4-large"`, `"voyage-code-4"`, ...	the Voyage API (needs `VOYAGE_API_KEY`)
`"hashing"`	a hashed bag of words without dependencies, for tests and quick baselines
an embedder object	used as it is (cached under its `name`; an object without a name is not cached)

"voyage-4-nano@512" (or embedding_kwargs={"dimension": 512}) truncates Matryoshka embeddings after the cache, so the full vector is computed once and serves every dimension.
Local models run in bfloat16 on CUDA (float32 elsewhere), with the token limit capped at the model's maximum.
Voyage API models (pymongopredict.embeddings_api.embed.VoyageEmbedder) take output_dimension (256, 512, 1024 or 2048), api_key, max_retries (default 5, with exponential backoff on rate limits) and timeout through embedding_kwargs. Texts go in requests of at most 1,000 texts and within the model's tokens per request. Contextualized and multimodal models (voyage-context-*, voyage-multimodal-*) are not supported.
mask_target=True keeps the target as the first key with the value <MASK> and puts Predict the value of <MASK>: before the text (mask_token and prefix change them), so the model knows which field is asked for. This is the input format of the voyage-json models, which are trained to predict from whole documents.

Examples

Usage:

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from pymongopredict import DocumentFeaturizer, MongoDataset

ds = MongoDataset(collection, target="room_type")
train, test = train_test_split(ds, test_size=0.2, stratify=ds.y, random_state=0)
prep = DocumentFeaturizer(columns=[...], embed_fields=[...])   # docs -> X
clf = LogisticRegression().fit(prep.fit_transform(train), train.y)
clf.score(prep.transform(test), test.y)

class MongoDataset:

Documents of a collection (or an explicit list) together with an optional target.

The query is lazy: nothing is read until the first access to len(), ids, y, iter_batches() or field_values() (or an explicit load()), and the result is cached.

Parameters

collection (pymongo Collection or None): Anything with .find (mongomock works). None for in-memory documents (see from_documents()).
target (str, optional): Dotted path of the field to predict (e.g. "imdb.rating"). Documents without a value are dropped. None (default) for an unlabelled dataset; y is then None.
filter, projection, limit, sort (optional): Forwarded to collection.find. The projection defines what a document looks like in both modes; the target field is added to an inclusion projection (or removed from the excluded fields of an exclusion projection) so y can always be read. limit of None or 0 means no limit.
streaming (bool, default False): If True the documents are never held in memory all at once. load() only reads _id and the target; documents are fetched again in _id batches whenever an estimator iterates over the dataset (with a projection limited to the fields it uses). Use it when documents are much larger than the features derived from them. The collection must not change in between: a document that disappeared raises KeyError.
batch_size (int, default 512): Documents per batch in iter_batches() (and per _id query in streaming mode).
prefetch (int, default 2): In streaming mode, number of batches fetched ahead in a background thread so the database round trip overlaps with featurisation / embedding (0 disables).

Raises

ValueError: If batch_size is smaller than 1.

Notes

Each document is identified by its _id. When a document has no _id (in-memory documents, or an in-memory projection with "_id": 0) its position is used instead.

Examples

>>> ds = MongoDataset(db.movies, target="imdb.rating", projection={"title": 1, "plot": 1})
>>> ds.y.dtype
dtype('float64')
>>> ds.field_values("title")[:2]

MongoDataset( collection, target: str | None = None, filter: dict | None = None, projection: dict | None = None, limit: int | None = None, sort=None, streaming: bool = False, batch_size: int = 512, prefetch: int = 2)

collection

target

filter

projection

limit

sort

streaming

batch_size

prefetch

@classmethod

def from_documents(cls, documents: Iterable[dict], target: str | None = None, y=None):

In-memory dataset from a list / cursor of documents.

Parameters

documents (iterable of dict): Documents (a list, a pymongo cursor, a generator). It is read once into a list.
target (str, optional): Dotted path of the label inside each document. Documents missing it are dropped. Ignored for labels when y is given, but still stored as target.
y (array-like, optional): Labels given separately, one per document (no document is dropped). Converted with _as_target_array().

Returns

MongoDataset: A dataset with no collection; it cannot be reloaded, and fields left out of the documents cannot be fetched later.

Raises

ValueError: If y is not one dimensional or its length differs from the number of documents.

Examples

>>> ds = MongoDataset.from_documents([{"a": 1, "t": "x"}, {"a": 2}], target="t")
>>> len(ds)
1
>>> ds = MongoDataset.from_documents([{"a": 1}, {"a": 2}], y=[0.5, 1.5])

def load(self, force: bool = False):

Run the query and cache the result.

In memory, the documents are fetched with the dataset's projection (extended to the target). When streaming, only _id and the target are fetched. Documents with a missing target are dropped. A reload clears the cache of field_values().

Parameters

force (bool, default False): Run the query again even if the dataset is already loaded (e.g. after the collection changed).

Returns

MongoDataset: self, so calls can be chained (ds.load().y).

Raises

ValueError: If the dataset is a subset (from subset() or a scikit-learn split), which only knows its ids, or if it has no collection and was not built with from_documents().

ids: list

_id of every document, in dataset order (a copy).

Documents without an _id (in-memory documents, or a projection with "_id": 0) are identified by their position in the query result instead.

y: numpy.ndarray | None

Target values aligned with ids (None when the dataset has no target).

Numeric targets come back as a float array, booleans as bool, anything else as strings, so the array can go straight into an sklearn estimator.

documents: list[dict]

All documents as a list, in dataset order (a new list on every access).

In streaming mode this fetches every document (defeating the purpose), so prefer iter_batches().

def iter_batches(self, projection: dict | None = None):

Yield lists of at most batch_size documents, in dataset order.

Parameters

projection (dict, optional): Fields the consumer needs, used to limit what is fetched in streaming mode. It is ignored for in-memory documents (already loaded) and for datasets that have their own projection, since that projection defines the documents in both modes.

Yields

list of dict: Consecutive batches. In streaming mode each batch is one _id query; with prefetch > 0 and more than one batch, the next batches are fetched in a background thread while the current one is processed.

Raises

KeyError: In streaming mode, if documents were deleted from the collection after loading.

Examples

>>> for batch in ds.iter_batches(projection={"title": 1}):
...     texts = [d.get("title") for d in batch]

def field_values(self, path: str) -> list:

Values of a dotted field for every document, aligned with ids.

Used for stratify or group values in scikit-learn splits, and by featurizers that read single columns. Streaming datasets fetch just that field (once, then cached until the next load()) with a projection, and so do in-memory datasets whose projection left the field out, so both modes see the same values.

Parameters

path (str): Dotted field path, e.g. "host.host_id" or "scores.0".

Returns

list: One value per document (a new list). Missing fields give None; BSON scalars (Decimal128, ObjectId) are converted to plain Python values.

Raises

ValueError: For an in-memory dataset whose projection leaves out both path and _id (or that has no collection), since the field then cannot be fetched by _id.

def subset( self, idx: Sequence[int]) -> MongoDataset:

Dataset restricted to positional indices.

The subset shares the collection, query settings and mode (in memory or streaming), and copies the matching ids, targets, documents and cached field_values(). It is a view: it cannot load() its query again.

Parameters

idx (sequence of int): Positions into this dataset (e.g. from a scikit-learn splitter). Order is kept and repeats are allowed.

Returns

MongoDataset: A new dataset with len(idx) documents.

Examples

>>> first_ten = ds.subset(range(10))

shape: tuple[int]

(n_documents,), as scikit-learn expects of an array-like (loads the dataset).

def take( self, indices, axis: int = 0) -> MongoDataset:

The documents at indices as a new dataset (subset()); axis must be 0.

This is what sklearn.model_selection.train_test_split and the other scikit-learn splitters call, so they split a dataset into datasets:

>>> from sklearn.model_selection import train_test_split
>>> train, test = train_test_split(ds, test_size=0.2, stratify=ds.y, random_state=0)

iloc

Positional indexing as in pandas: ds.iloc[10:20] or ds.iloc[[0, 5]].

class MongoSeries:

A univariate time series aggregated from the documents of a collection.

Only the two fields are read, streamed through the cursor. series is a pd.Series with a PeriodIndex, ready for sktime. The collection is read lazily, on the first access to series, len() or split(), and the result is cached.

Parameters

collection (pymongo.collection.Collection): The collection to read.
time_field (str): Dotted path of the timestamp field. Accepted values: datetime / date / pd.Timestamp, ISO 8601 strings, and numbers as epoch seconds (or epoch milliseconds when above 1e11). Timezone-aware times are converted to naive UTC. Documents whose time cannot be parsed are skipped.
value_field (str): Dotted path of the numeric field to aggregate. Decimal128 values are read as floats; values that are missing, None or not convertible to float are skipped.
freq (str, default "D"): pandas offset alias used to bucket timestamps ("h", "D", "W", "MS" ...), with an optional multiple and anchor ("2MS", "QS-FEB"). The period-style aliases "M", "Q", "Y" and "A" are accepted and resampled as "ME", "QE" and "YE", which pandas 3 requires. Month, quarter and year aliases (also business ones) map to the matching periods, keeping the multiple ("2MS" gives "2M" periods) and ending start-anchored periods the month before the anchor ("QS-FEB" gives "Q-JAN"). For other aliases the period frequency is inferred from the resampled index.
agg (str, default "mean"): Aggregation within a bucket, any name pandas Resampler.agg accepts ("mean", "sum", "count", "median", "max" ...).
fill ({"interpolate", "ffill"}, number or None, default "interpolate"): How to fill empty buckets: "interpolate" interpolates linearly and extends the first and last values to the ends; "ffill" carries the last value forward (and the first value backward for leading gaps); a number fills with that value; None leaves NaN. Aggregations such as "count" and "sum" give 0, not NaN, for empty buckets, so there is nothing to fill.
filter (dict or None, default None): find filter; None reads every document.
valid (callable or None, default None): valid(value) -> bool to drop sentinel values (e.g. 999.9 in the weather sample data). It sees the value before conversion to float.
limit (int or None, default None): Read at most this many documents (in natural order); None or 0 reads all.

Attributes

series (pandas.Series): The aggregated series, named after value_field.

Examples

>>> ms = MongoSeries(coll, "ts", "airTemperature.value", freq="D")  # doctest: +SKIP
>>> train, test = ms.split(0.2)  # doctest: +SKIP

MongoSeries( collection, time_field: str, value_field: str, freq: str = 'D', agg: str = 'mean', fill='interpolate', filter: dict | None = None, valid=None, limit: int | None = None)

collection

time_field

value_field

freq

agg

fill

filter

valid

limit

def load(self, force: bool = False):

Read the collection and build the series (once, unless force).

Parameters

force (bool, default False): Re-read the collection even if the series is already cached, e.g. after new documents were inserted.

Returns

MongoSeries: self, so that ms.load().series chains.

Raises

ValueError: If fewer than three documents have a usable (time, value) pair, or fewer than three periods remain after resampling.

series: pandas.Series

The aggregated series (float values, PeriodIndex); loads it on first access.

def split( self, test_size: float | int = 0.2) -> tuple[pandas.Series, pandas.Series]:

Trailing test_size periods (int) or fraction (float) become the test series.

Parameters

test_size (float or int, default 0.2): An integer of at least 1 is a number of periods; anything else is a fraction of the series length (rounded). The test size is clipped so that the test series has at least one period and the training series at least two.

Returns

tuple of pandas.Series: (train, test), contiguous and in time order.

Examples

>>> train, test = ms.split(test_size=12)  # last 12 periods  # doctest: +SKIP

class DocumentFeaturizer(sklearn.base.BaseEstimator, sklearn.base.TransformerMixin):

Turn documents into [tabular columns | embedding of JSON] features.

A scikit-learn transformer. All constructor arguments are stored unchanged (so get_params, clone and grid search work); everything learned is set by fit in attributes ending with _, and every fit starts from scratch.

Parameters

columns (list of str or None, default None): Dotted field paths ("imdb.rating", "tags.0") used directly as tabular features. None or [] means no tabular block. The kind of each column is inferred from the training documents (see _infer_kind):
- "size": every value is a list or sub-document; the feature is its length.
- "bool": every value is a boolean; at transform time 0 / 1, "yes" / "no", "true" / "false", "t" / "f" and "y" / "n" are read too, anything else is missing.
- "datetime": datetime values; the feature is seconds since 1970-01-01 UTC.
- "numeric": a numeric dtype, or at least 95% of the values parse as numbers (numeric strings included). Unparsable values and infinities are missing, values beyond the float32 range are clipped. A column with no values at all is numeric.
- "categorical": anything else. Values become strings (1, 1.0 and "1" are one category) and missing values the category "<NA>".
Non-categorical columns are imputed with the training median (and standardised when scale_numeric); categorical columns are one-hot encoded.
embed_fields (list of str, bool or None, default None): What is serialised to JSON and embedded. None or True: the whole document except _id and the excluded paths (the columns stay in the JSON too). A list of dotted paths: only those fields, in document order. False or []: no embedding block.
exclude (sequence of str, default ()): Dotted paths never embedded. The target of a MongoDataset passed to fit is excluded automatically; exclude_ holds the effective list after fitting.
embedding_model (str, BaseEmbedder or None, default "voyage-4-nano"): Embedding model: "voyage-4-nano" or any "org/name" Hugging Face id (local), other "voyage-*" names (Voyage API), "hashing" (tests), a "@<dim>" suffix for Matryoshka truncation, or an embedder instance.
embedding_kwargs (dict or None, default None): Forwarded to the embedder constructor (e.g. {"dimension": 512}, {"batch_size": 64}).
cache_dir (str or None, default ".mongopredict_cache"): Directory of the sqlite embedding cache, keyed by the embedder's settings and the text; None disables it.
batch_size (int, default 512): Documents per batch for cursors, generators and lists. A MongoDataset uses its own batch_size.
mask_target (bool, default False): False removes the target field from the embedded JSON. True keeps it as the first key, with the value mask_token, and puts prefix before the JSON, so the model sees which field is asked for. The masked key is mask_key or, failing that, the dataset target; with plain documents (no dataset target) only mask_key can be masked, and it is then also excluded from the rest of the JSON.
mask_token (str, default "<MASK>"): The value shown for the masked field.
mask_key (str or None, default None): The key shown with the mask (default: the dataset target). Useful when the target is a derived field (e.g. a collapsed or logged copy) and the model should see the original field name.
prefix (str or None, default None): Text put before the JSON. None means "Predict the value of <mask_token>: " when a field is masked and nothing otherwise; any string ("" included) is used as is.
max_json_chars (int or None, default 8000): Character budget for the whole text (prefix included). Fields are added in order and the first one that does not fit is cut (when more than its key fits), the rest dropped. None means unlimited. Ignored when max_tokens is set.
max_tokens (int, "model" or None, default None): Token budget instead of characters, measured with tokenizer. "model" uses the tokenizer's model_max_length, capped by the model config's max_position_embeddings (32768 for voyage-4-nano). 16 tokens are kept free for the document prompt and special tokens. For a local embedder the model's max_seq_length is set to this budget too, unless embedding_kwargs sets it.
tokenizer (str, tokenizer object or None, default None): Used only to measure max_tokens. A Hugging Face id, or a Hugging Face style tokenizer (callable on a list of texts, returning {"input_ids": ...}, with decode). None: the embedding model's tokenizer when it is a local model, else voyage-4-nano's.
scale_numeric (bool, default True): Standardise the non-categorical columns after imputation.
min_frequency (int, default 5): Categories seen fewer times in training go to the one-hot "infrequent" column.
max_categories (int, default 100): Upper bound on one-hot columns per categorical column (infrequent bucket included). The feature matrix is dense float32: its memory is documents x (embedding dimension + numeric columns + one-hot columns) x 4 bytes, so many categorical columns with a large max_categories cost memory on large training sets.

Attributes

target_ (str or None): The dataset target seen by the last fit (None for plain documents).
exclude_ (list of str): Paths left out of the embedded JSON: exclude, the target and, for plain documents in mask mode, mask_key.
embedder_ (BaseEmbedder or None): The embedder built by the last fit (None when embeddings are disabled).
tokenizer_ (tokenizer or None): Tokenizer loaded for max_tokens; reset by every fit.
column_kinds_ (dict of str to str): Inferred kind of each column (see columns).
tabular_ (sklearn ColumnTransformer or None): Fitted preprocessing of the columns (None without columns).
n_tabular_features_, n_embedding_features_ (int): Widths of the two blocks of the output matrix.

Examples

>>> docs = [{"_id": i, "price": 10 + i, "kind": "a" if i % 2 else "b",
...          "text": f"item {i}"} for i in range(20)]
>>> prep = DocumentFeaturizer(columns=["price", "kind"], embedding_model="hashing",
...                           embedding_kwargs={"n_features": 32}, cache_dir=None)
>>> X = prep.fit_transform(docs)  # the whole document but ``_id`` is embedded
>>> X.shape, X.dtype
((20, 35), dtype('float32'))
>>> list(prep.get_feature_names_out()[:4])
['num__price', 'cat__kind_a', 'cat__kind_b', 'emb_0']

DocumentFeaturizer( columns: Optional[Sequence[str]] = None, embed_fields: Union[Sequence[str], NoneType, bool] = None, exclude: Sequence[str] = (), embedding_model: str | BaseEmbedder | None = 'voyage-4-nano', embedding_kwargs: dict | None = None, cache_dir: str | None = '.mongopredict_cache', batch_size: int = 512, mask_target: bool = False, mask_token: str = '<MASK>', mask_key: str | None = None, prefix: str | None = None, max_json_chars: int | None = 8000, max_tokens: int | str | None = None, tokenizer=None, scale_numeric: bool = True, min_frequency: int = 5, max_categories: int = 100)

columns

embed_fields

exclude

embedding_model

embedding_kwargs

cache_dir

batch_size

mask_target

mask_token

mask_key

prefix

max_json_chars

max_tokens

tokenizer

scale_numeric

min_frequency

max_categories

def token_limit(self) -> int | None:

Effective token budget per text.

Returns

int or None: max_tokens as an int; for max_tokens="model" the tokenizer's model_max_length capped by the model's max_position_embeddings. None when max_tokens is None (the character budget max_json_chars is used).

Raises

ValueError: With max_tokens="model" when the tokenizer reports no usable length.

def fit(self, X, y=None):

Fit the featurizer (column kinds, tabular preprocessing, embedder).

Parameters

X (MongoDataset, cursor, iterable of dict or list of dict): Training documents. A MongoDataset is read in its own batches and its target is excluded from (or masked in) the embedded JSON. Other inputs are read in batches of batch_size.
y (ignored): Present for sklearn compatibility; take labels from MongoDataset.y.

Returns

DocumentFeaturizer: self. The documents are embedded during the fit, so with a cache a following transform of the same documents is cheap.

Raises

TypeError: If columns, embed_fields or exclude is a single string.
ValueError: If there are no documents, or neither columns nor embeddings.

def fit_transform(self, X, y=None, **fit_params):

Fit on X and return its feature matrix, reading the documents once.

Parameters

X (MongoDataset, cursor, iterable of dict or list of dict): As in fit().
y (ignored):
**fit_params (ignored):

Returns

numpy.ndarray of float32, shape (n_documents, n_features): [tabular columns | embedding]; column names from get_feature_names_out().

def transform(self, X):

Feature matrix of X with the fitted preprocessing and embedder.

Parameters

X (MongoDataset, cursor, iterable of dict or list of dict): Documents to featurize. The target is excluded (or masked) as in the fit, so it may be present or absent.

Returns

numpy.ndarray of float32, shape (n_documents, n_features): An empty (0, n_features) array when X has no documents.

def fit_batches(self, batches):

Fit from an iterable of document batches (lists of dicts) without holding them all.

Batches carry no target, so nothing is excluded automatically: pass the target in exclude (or set mask_key with mask_target=True). Returns self.

def fit_transform_batches(self, batches) -> numpy.ndarray:

Fit and return the feature matrix, consuming the document batches once.

Only the values of columns and the embeddings are kept per batch, so memory stays proportional to the features rather than to the raw documents. As in fit_batches(), no target is known, so it must be in exclude.

def transform_batches(self, batches) -> numpy.ndarray:

Transform an iterable of document batches.

Embeddings are computed batch by batch (so a streaming dataset is never held in memory), while the column values, which are small, are collected and encoded in a single sklearn call at the end: per batch sklearn overhead dominated on collections of small documents.

Parameters

batches (iterable of list of dict): Document batches, e.g. from MongoDataset.iter_batches.

Returns

numpy.ndarray of float32, shape (n_documents, n_features)

def get_feature_names_out(self, input_features=None):

Names of the output columns, in order.

Parameters

input_features (ignored): Present for sklearn compatibility.

Returns

numpy.ndarray of str (object dtype): The tabular names from the ColumnTransformer (e.g. "num__year", "cat__rated_PG", "cat__rated_infrequent_sklearn"), then "emb_0" ... "emb_<d-1>" for the embedding dimensions.

class SentenceTransformerEmbedder(pymongopredict.embeddings_api.BaseEmbedder):

Local embeddings through sentence-transformers.

The model is loaded lazily on first use (see model) and is never pickled, so a fitted featurizer can be saved and loaded on another host.

Parameters

model_name (str, default "voyageai/voyage-4-nano"): Hugging Face id, or a key of LOCAL_ALIASES ("voyage-4-nano"). voyage-4-nano is Apache 2.0, 346M parameters.
batch_size (int, default 32): Maximum texts per forward pass.
max_seq_length (int or None, default 1024): Token limit per text; longer JSON documents are truncated. Capped at the model's own limit (e.g. 256 for all-MiniLM-L6-v2). None keeps the model's limit.
device (str or None, default None): "cuda", "cpu" (or e.g. "cuda:1"); None picks CUDA when available.
normalize (bool, default True): Return unit-length vectors.
trust_remote_code (bool, default True): Forwarded to SentenceTransformer (needed by models with custom code).
show_progress_bar (bool, default False): Forwarded to encode.
dtype (str, default "auto"): "auto" runs in bfloat16 on CUDA devices that support it and float32 elsewhere; or an explicit torch dtype name such as "float32" / "float16". float16 overflows in voyage-4-nano (Qwen3 backbone) and yields NaN vectors: if any vector of a call comes back non-finite, a RuntimeWarning is issued, the model is switched to float32 for good and those rows are recomputed.
max_batch_tokens (int, default 65536): Texts are grouped by length so that a batch holds at most about this many tokens (estimated as characters / 3, capped by max_seq_length). This keeps long documents from exhausting GPU memory while short ones still go through in batches of batch_size. On a CUDA out-of-memory error a group is retried with a quarter of the batch size, down to one text, before the error is raised.
prompt (str or None, default None): Text put before every document instead of the model's document prompt. None: the model's own (e.g. voyage-4-nano's "Represent the document for retrieval: "); "": none.
max_seq_length_in_key (bool, default True): Whether |len=<max_seq_length> goes into name (the cache key). The featurizer turns it off when it has already cut every text to the model's length with the model's own tokenizer, so those vectors are the same as untruncated ones.

Attributes

name (str): Cache key prefix: the model id, then |prompt=<prompt> for a custom prompt, |raw when not normalised, and |len=<max_seq_length> (see max_seq_length_in_key).
part (tuple of (int, int) or None): Set only by _CachedEmbedder during one call in part mode.
last_done (numpy.ndarray of bool or None): Rows the last embed() call computed (all of them outside part mode).

Notes

Matryoshka truncation is not done here: pass dimension= to _make_embedder(), which caches the full vectors once and truncates on the way out (see _MatryoshkaTruncate).

SentenceTransformerEmbedder( model_name: str = 'voyageai/voyage-4-nano', batch_size: int = 32, max_seq_length: int | None = 1024, device: str | None = None, normalize: bool = True, trust_remote_code: bool = True, show_progress_bar: bool = False, dtype: str = 'auto', max_batch_tokens: int = 65536, prompt: str | None = None, max_seq_length_in_key: bool = True)

model_name

prompt

batch_size

max_batch_tokens

max_seq_length

device

normalize

trust_remote_code

show_progress_bar

dtype

name = 'base'

model

The SentenceTransformer, loaded on first access.

Resolves device and dtype="auto" at load time and sets the model's max_seq_length to min(max_seq_length, native limit).

def embed(self, texts: Sequence[str]) -> numpy.ndarray:

Embed texts in length-grouped batches.

Parameters

texts (sequence of str):

Returns

numpy.ndarray of float32, shape (len(texts), dim): Rows in the order of texts. In part mode (part set) only this part's rows are computed; the others are zero and last_done marks which rows are real.

part: tuple[int, int] | None = None

last_done = None

def part_of(self, groups, texts) -> list[int]:

Assign every length group to one of the K parts of part = (k, K).

Each group goes to the least loaded part, the load being the group's padded tokens (estimated as characters / 3), in the groups' order (longest first). The assignment is deterministic and the loads carry over between calls, so the K processes, each running the same calls, agree on who computes which group.

Parameters

groups (list of list of int): Groups from _token_groups.
texts (list of str): The texts the indices refer to.

Returns

list of int: The part (0 .. K-1) of each group.

class VoyageEmbedder(pymongopredict.embeddings_api.BaseEmbedder):

Embeddings through the Voyage AI API (the voyageai client is a dependency of the package; needs an API key in VOYAGE_API_KEY or api_key).

Any text embedding model the API serves works, e.g. "voyage-4-lite", "voyage-4", "voyage-4-large", "voyage-code-4", "voyage-finance-2" or "voyage-law-2". voyage-4-nano is an open-weight model and always runs locally (see SentenceTransformerEmbedder). Texts are sent in requests of at most batch_size texts (the API takes up to 1,000) and at most the model's tokens per request (1M for voyage-4-lite, 320K for voyage-4, 120K for voyage-4-large and the others), estimated as characters / 3. Rate limits and server errors are retried by the client, with exponential backoff, up to max_retries times.

Parameters

model (str, default "voyage-4-lite"): Voyage model name.
output_dimension (int or None, default None): Dimension requested from the API (256, 512, 1024 or 2048 for the voyage-4 models); None is the model's default (1024).
input_type (str or None, default "document"): Voyage input_type ("document", "query" or None).
batch_size (int, default 1000): Most texts per API request.
api_key (str or None, default None): API key; None reads VOYAGE_API_KEY.
truncation (bool, default True): Let the API truncate texts longer than the model's context (32,000 tokens for the voyage-4 models) instead of failing.
max_retries (int, default 5): Retries of a failed request (rate limit, timeout, server error), passed to voyageai.Client.
timeout (float or None, default None): Seconds before a request times out; None is the client's default.

Attributes

name (str): "voyage-api:<model>@<output_dimension or full>", plus |<input_type>|truncation=... when those differ from the defaults (so default keys stay as they were).

Raises

ValueError: For a model served by another endpoint (voyage-context-*, voyage-multimodal-*), and, on the first request, when no API key is set.
ImportError: On the first request, when the voyageai package is not installed.

VoyageEmbedder( model: str = 'voyage-4-lite', output_dimension: int | None = None, input_type: str | None = 'document', batch_size: int = 1000, api_key: str | None = None, truncation: bool = True, max_retries: int = 5, timeout: float | None = None)

model

output_dimension

input_type

batch_size

api_key

truncation

max_retries

timeout

name = 'base'

client

The voyageai.Client, created on first use (so importing needs no voyageai).

def embed(self, texts: Sequence[str]) -> numpy.ndarray:

Embed texts with as few API requests as the limits allow.

Returns

numpy.ndarray of float32, shape (len(texts), dim)

class HashingEmbedder(pymongopredict.embeddings_api.BaseEmbedder):

Hashed word uni- and bigram vectors: fast, deterministic and dependency free.

Not a semantic embedding; use it for tests, CI, quick pipeline checks and as a cheap bag-of-tokens baseline. analyzer="char_wb" with ngram_range=(3, 5) gives character n-grams instead (about 10x slower on long JSON documents).

Parameters

n_features (int, default 256): Output dimension (number of hash buckets).
ngram_range (tuple of (int, int), default (1, 2)): Range of n-gram lengths.
analyzer (str, default "word"): HashingVectorizer analyzer ("word", "char" or "char_wb").

Attributes

name (str): "hashing@<n_features>", plus |<analyzer><ngram_range> for non-default analyzers, so different settings never share cache entries.

Examples

>>> HashingEmbedder(n_features=8).embed(['{"a": 1}']).shape
(1, 8)

HashingEmbedder(n_features: int = 256, ngram_range=(1, 2), analyzer: str = 'word')

n_features

name = 'base'

def embed(self, texts: Sequence[str]) -> numpy.ndarray:

L2-normalised hashed n-gram counts of texts, shape (len(texts), n_features).

class BaseEmbedder:

Interface of all embedders: embed(texts) -> np.ndarray of shape (n_texts, dim).

Subclass it to plug in any model: implement embed() and set name to a string that identifies the model and every setting that changes its vectors. name is the cache key prefix of _CachedEmbedder; an embedder that keeps the default "base" is not cached by _make_embedder().

Attributes

name (str): Identifier of the model and its settings, used in cache keys.

name: str = 'base'

def embed(self, texts: Sequence[str]) -> numpy.ndarray:

Embed texts.

Parameters

texts (sequence of str):

Returns

numpy.ndarray of float32, shape (len(texts), dim)