LemurTrainer

Trains a LEMUR fixed dimensional encoder for a late interaction model and corpus. Training is a separate pipeline step; the saved artifact is then loaded through the embeddings vectors.lemur configuration.

Example

from txtai.pipeline import LemurTrainer

corpus = [
    "First document",
    "Second document",
    "Third document",
]

trainer = LemurTrainer()
trainer(
    "colbert-ir/colbertv2.0",
    corpus,
    "lemur-model",
    epochs=100,
    validation_split=0.1,
)

epochs is required so the training mode is explicit. Use epochs=100 for the quality-oriented MLP setting. This can take hours on a CPU. Use epochs=0 to select deterministic random ELM features as a lower-cost fallback.

The learn distribution defaults to learn_category="query". Target documents are always encoded with the data encoder, so the default encodes selected corpus texts once as data and once as queries. Set learn_category="data" to reuse the data encodings, or pass a separate iterable of texts with learn.

Set corpus_subset_size to a positive integer to sample that many raw corpus texts under seed before either encoding pass. The default is None, which uses every corpus text. train_subset_size and learn_subset_size apply later, after token vectors have already been created. corpus_subset_size applies to data; a separate learn iterable remains caller-sized.

For trained MLP features, set validation_split to a fraction greater than zero and less than one to retain the epoch with the lowest held-out loss. The default is 0.0, which preserves training-loss selection. The selected one-based epoch, loss and metric are available as selected_epoch, selected_loss and selection_metric on the fitted or reloaded encoder.

The artifact contains config.json and model.safetensors. It stores the inference feature model, output-normalization statistics and token sample needed to encode documents added after training. The training-only output readout is not saved.

Load the artifact in an embeddings configuration.

embeddings:
    path: colbert-ir/colbertv2.0
    vectors:
        lemur:
            path: lemur-model

Search behavior

LEMUR approximates standardized MaxSim targets, so useful ranking scores can be negative. txtai's dense vector path L2-normalizes the fixed vectors and removes results with scores less than or equal to zero. This can change ordering and return fewer than the requested number of candidates compared with raw maximum inner-product search.

The Faiss backend uses exact IDMap,Flat search through 5,000 rows and switches to an IVF index above that threshold. In the measured scifact run, default IVF reduced LEMUR NDCG@10 by 43% relative to exact search, compared with 25% for MUVERA. For LEMUR corpora above 5,000 rows, either pin faiss.components to an exact index or tune the IVF settings for the corpus. The exact configuration is:

embeddings:
    path: colbert-ir/colbertv2.0
    vectors:
        lemur:
            path: lemur-model
    faiss:
        components: IDMap,Flat

Methods

Python documentation for the pipeline.

__call__(path, data, output, gpu=True, method=None, tokenizer=None, maxlength=None, vectors=None, learn=None, learn_category='query', corpus_subset_size=None, validation_split=0.0, **kwargs)

Trains a LEMUR feature encoder.

Parameters:

Name	Type	Description	Default
path		late interaction model path	required
data		iterable of corpus texts	required
output		artifact output directory	required
gpu		tensor accelerator setting	True
method		optional pooling method	None
tokenizer		optional tokenizer path	None
maxlength		maximum token length	None
vectors		additional model arguments	None
learn		optional iterable of texts used to learn features, defaults to data	None
learn_category		encoder category for learn texts (data or query)	'query'
corpus_subset_size		optional maximum number of corpus texts selected before encoding	None
validation_split		fraction of sampled learn tokens held out for validation	0.0
kwargs		LEMUR fit arguments	{}

Returns:

Type	Description
	fitted LEMUR encoder

Source code in txtai/pipeline/train/lemur.py

def __call__(
    self,
    path,
    data,
    output,
    gpu=True,
    method=None,
    tokenizer=None,
    maxlength=None,
    vectors=None,
    learn=None,
    learn_category="query",
    corpus_subset_size=None,
    validation_split=0.0,
    **kwargs,
):
    """
    Trains a LEMUR feature encoder.

    Args:
        path: late interaction model path
        data: iterable of corpus texts
        output: artifact output directory
        gpu: tensor accelerator setting
        method: optional pooling method
        tokenizer: optional tokenizer path
        maxlength: maximum token length
        vectors: additional model arguments
        learn: optional iterable of texts used to learn features, defaults to data
        learn_category: encoder category for learn texts (data or query)
        corpus_subset_size: optional maximum number of corpus texts selected before encoding
        validation_split: fraction of sampled learn tokens held out for validation
        kwargs: LEMUR fit arguments

    Returns:
        fitted LEMUR encoder
    """

    data = list(data)
    if not data:
        raise ValueError("data must contain at least one corpus text")
    if learn_category not in ("data", "query"):
        raise ValueError("learn_category must be data or query")
    if kwargs.get("epochs") is None:
        raise ValueError("epochs must be set explicitly: use epochs=100 for trained MLP quality or epochs=0 for deterministic ELM features")
    if corpus_subset_size is not None:
        if isinstance(corpus_subset_size, bool) or not isinstance(corpus_subset_size, int) or corpus_subset_size <= 0:
            raise ValueError("corpus_subset_size must be a positive integer")
        if corpus_subset_size < len(data):
            indices = sorted(random.Random(kwargs.get("seed", 42)).sample(range(len(data)), corpus_subset_size))
            data = [data[index] for index in indices]

    learn = list(learn) if learn is not None else None
    if learn is not None and not learn:
        raise ValueError("learn must contain at least one text")

    deviceid = Models.deviceid(gpu)
    modelargs = {**(vectors if vectors else {}), **{"muvera": None, "lemur": None}}
    pooling = PoolingFactory.create(
        {
            "method": method,
            "path": path,
            "device": deviceid,
            "tokenizer": tokenizer,
            "maxlength": maxlength,
            "modelargs": modelargs,
        }
    )

    # A batch size of one preserves each document's true token count. The late
    # pooling path normalizes token rows before returning raw multi-vectors.
    documents = [pooling.encode([text], batch=1, category="data")[0] for text in data]
    learn_documents = None
    if learn is not None or learn_category != "data":
        learn = data if learn is None else learn
        learn_documents = [pooling.encode([text], batch=1, category=learn_category)[0] for text in learn]

    lemur = Lemur(device=Models.device(deviceid))
    return lemur.fit(documents, output=output, validation_split=validation_split, learn=learn_documents, **kwargs)