HuggingFace Adapter

The [huggingface] extra provides two FAISS-indexed retrievers backed by HuggingFace datasets.

pip install "schematize[huggingface]"

Both retrievers:

Load a HuggingFace dataset and embed it with a sentence-transformers model
Build a FAISS index lazily on first call
Cache the index to disk — subsequent instantiations load from cache instead of re-embedding

HuggingFaceRetriever

General-purpose retriever. You choose the dataset, the text column, and the embedding model.

from schematize.retrieval.huggingface import HuggingFaceRetriever

retriever = HuggingFaceRetriever(
    dataset_name="JuDDGES/pl-court-raw",
    text_column="text",
    embedding_model="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    max_documents=50_000,          # None = entire dataset
    split="train",
    batch_size=256,
    device=None,                   # None = auto-detect (GPU if available)
    index_path=".cache/my-index",  # omit to disable disk caching
)

Index caching

The first call builds and embeds the dataset (can take a while for large datasets). With index_path set, the FAISS index and dataset arrow files are saved to disk. Subsequent instantiations detect the cache, validate metadata (dataset name, column, model, split, max_documents), and load from disk if all match.

To force a rebuild:

import asyncio
asyncio.run(retriever.build_index())

Rebuilding with different settings

Change any constructor parameter and the metadata check will detect the mismatch — the cache is ignored and the index is rebuilt.

MMLWRobertaV2Retriever

A Polish-optimised retriever that doubles as a worked example of a custom retriever — it specialises the HuggingFace base for a specific model and language, which makes it a useful template for building your own (see the custom retriever guide).

It uses sdadas/mmlw-retrieval-roberta-large-v2, a model fine-tuned for retrieval on Polish benchmarks (NDCG@10 of 60.71 on PIRB). Queries are automatically prefixed with [query]: as required by the model.

from schematize.retrieval.huggingface import MMLWRobertaV2Retriever

retriever = MMLWRobertaV2Retriever(
    dataset_name="JuDDGES/pl-court-raw",
    text_column="text",
    max_documents=50_000,
    index_path=".cache/mmlw-index",
)

The model is loaded in bfloat16 by default (as recommended by the model authors). Override with model_kwargs:

retriever = MMLWRobertaV2Retriever(
    dataset_name="JuDDGES/pl-court-raw",
    text_column="text",
    model_kwargs={"model_kwargs": {"torch_dtype": "float32"}},
)

Choosing the right retriever

Use HuggingFaceRetriever for any language or model. Use MMLWRobertaV2Retriever if you're working in Polish — it's pre-configured for sdadas/mmlw-retrieval-roberta-large-v2, applies the [query]: prefix automatically, and uses about 1.3 GB in bfloat16.

API reference

HuggingFaceRetriever

HuggingFaceRetriever(
    dataset_name: str,
    text_column: str,
    embedding_model: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    max_documents: int | None = None,
    split: str = "train",
    batch_size: int = 256,
    device: str | None = None,
    index_path: str | Path | None = None,
    model_kwargs: dict | None = None,
)

Bases: _BaseHuggingFaceRetriever

DocumentRetriever backed by a HuggingFace dataset + FAISS index.

Embeddings are computed with a sentence-transformers model. The index is built lazily on first use. Pass index_path to persist the index to disk — subsequent instantiations will load it instead of re-embedding.

Parameters:

Name	Type	Description	Default
`dataset_name`	`str`	HuggingFace dataset identifier (e.g. `"JuDDGES/pl-court-raw"`).	required
`text_column`	`str`	Dataset column to embed and return as search target.	required
`embedding_model`	`str`	Sentence-transformers model name or path.	`'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'`
`max_documents`	`int \| None`	Number of rows to stream and index (`None` = no limit).	`None`
`split`	`str`	Dataset split to load.	`'train'`
`batch_size`	`int`	Batch size for embedding computation.	`256`
`device`	`str \| None`	Device for sentence-transformers (`None` = auto-detect).	`None`
`index_path`	`str \| Path \| None`	Directory to save/load the FAISS index and dataset arrow files.	`None`

MMLWRobertaV2Retriever

MMLWRobertaV2Retriever(
    dataset_name: str,
    text_column: str,
    max_documents: int | None = None,
    split: str = "train",
    batch_size: int = 256,
    device: str | None = None,
    index_path: str | Path | None = None,
    model_kwargs: dict | None = None,
)

Bases: _BaseHuggingFaceRetriever

DocumentRetriever using sdadas/mmlw-retrieval-roberta-large-v2 for Polish retrieval.

Optimised for information retrieval (NDCG@10 of 60.71 on PIRB). Queries are prefixed with [query]: as required by the model. Pass index_path to persist the index to disk — subsequent instantiations will load it instead of re-embedding.

Parameters:

Name	Type	Description	Default
`dataset_name`	`str`	HuggingFace dataset identifier.	required
`text_column`	`str`	Dataset column to embed and return as search target.	required
`max_documents`	`int \| None`	Number of rows to stream and index (`None` = no limit).	`None`
`split`	`str`	Dataset split to load.	`'train'`
`batch_size`	`int`	Batch size for embedding computation.	`256`
`device`	`str \| None`	Device for sentence-transformers (`None` = auto-detect).	`None`
`index_path`	`str \| Path \| None`	Directory to save/load the FAISS index and dataset arrow files.	`None`
`model_kwargs`	`dict \| None`	Extra kwargs forwarded to `SentenceTransformer`. Defaults to `{"model_kwargs": {"torch_dtype": "bfloat16"}}` as recommended.	`None`