XTRModel

class lightning_ir.models.xtr.model.XTRModel(config: XTRConfig, *args, **kwargs)[source]

Bases: ColModel

__init__(config: XTRConfig, *args, **kwargs) → None[source]

Methods

__init__(config, *args, **kwargs)

add_mask_scoring_input_ids()

doc_scoring_mask(input_ids, attention_mask)

encode_doc(encoding)

encode_query(encoding)

forward(query_encoding, doc_encoding[, num_docs])

from_pretrained(model_name_or_path, *args, ...)

from_xtr_checkpoint(model_name_or_path)

get_output_embeddings()

query_scoring_mask(input_ids, attention_mask)

score(query_embeddings, doc_embeddings[, ...])

Attributes

ALLOW_SUB_BATCHING

Flag to allow mini batches of documents for a single query.