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.

LemurTrainer owns training-data construction, feature-model setup, epoch training and validation-based selection. The loaded LEMUR pooler remains responsible for query-feature and document-weight encoding.

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,
    validationsplit=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 learncategory="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 learncategory="data" to reuse the data encodings, or pass a separate iterable of texts with learn.

When vectors.center is omitted, the trainer encodes the corpus without centering, computes one mean over all corpus token rows and uses that mean to center both data and learn vectors before fitting. The collection mean is stored in the LEMUR artifact. An explicit vectors.center setting keeps its configured behavior and does not store an automatic collection mean.

Set corpussubsetsize 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. trainsubsetsize and learnsubsetsize apply later, after token vectors have already been created. corpussubsetsize applies to data; a separate learn iterable remains caller-sized.

For trained MLP features, set validationsplit 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 selectedepoch, selectedloss and selectionmetric on the fitted or reloaded encoder.

MLP training displays a tqdm progress bar on an interactive terminal, including percent complete and the current validation loss when validationsplit is enabled (or training loss otherwise). Progress output is disabled for non-interactive runs.

The artifact contains config.json and model.safetensors. It stores the inference feature model, output-normalization statistics, token sample needed to encode documents added after training and, by default, the collection token mean. 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

Loading a LEMUR artifact with a stored collection mean automatically centers both query and document token vectors with that mean. This keeps training and search on the same representation and makes fixed vectors independent of indexing or query batch size. An explicit vectors.center setting takes precedence. Older artifacts without a stored mean use the standard late-pooling fallback: batch centering for models with multiple linear layers and no centering for models with zero or one linear layer.

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, learncategory='query', corpussubsetsize=None, validationsplit=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
learncategory		encoder category for learn texts (data or query)	'query'
corpussubsetsize		optional maximum number of corpus texts selected before encoding	None
validationsplit		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,
    learncategory="query",
    corpussubsetsize=None,
    validationsplit=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
        learncategory: encoder category for learn texts (data or query)
        corpussubsetsize: optional maximum number of corpus texts selected before encoding
        validationsplit: 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 learncategory not in ("data", "query"):
        raise ValueError("learncategory 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 corpussubsetsize is not None:
        if isinstance(corpussubsetsize, bool) or not isinstance(corpussubsetsize, int) or corpussubsetsize <= 0:
            raise ValueError("corpussubsetsize must be a positive integer")
        if corpussubsetsize < len(data):
            indices = sorted(random.Random(kwargs.get("seed", 42)).sample(range(len(data)), corpussubsetsize))
            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}}
    centerconfigured = vectors is not None and "center" in vectors
    if not centerconfigured:
        modelargs["center"] = False
    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]
    learndocuments = None
    if learn is not None or learncategory != "data":
        learn = data if learn is None else learn
        learndocuments = [pooling.encode([text], batch=1, category=learncategory)[0] for text in learn]

    centermean = None
    if not centerconfigured:
        centermean = np.concatenate(documents).mean(axis=0)
        pooling.center = {"scope": "collection", "mean": centermean}
        documents = pooling.centerdata(documents)
        if learndocuments is not None:
            learndocuments = pooling.centerdata(learndocuments)

    return self.fit(
        documents,
        output=output,
        device=Models.device(deviceid),
        validationsplit=validationsplit,
        learn=learndocuments,
        centermean=centermean,
        **kwargs,
    )