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)

encode(encoding[, expansion, ...])

Encodes a batched tokenized text sequences and returns the embeddings and scoring mask.

encode_doc(encoding)

Encodes tokenized documents.

encode_query(encoding)

Encodes tokenized queries.

forward(query_encoding, doc_encoding[, num_docs])

Embeds queries and/or documents and computes relevance scores between them if both are provided.

from_pretrained(model_name_or_path, *args, ...)

from_xtr_checkpoint(model_name_or_path)

get_output_embeddings()

Returns the output embeddings of the model for tieing the input and output embeddings.

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

Compute relevance scores between queries and documents.

scoring_mask(encoding[, expansion, ...])

Computes a scoring for batched tokenized text sequences which is used in the scoring function to mask out vectors during scoring.

Attributes

ALLOW_SUB_BATCHING

Flag to allow mini batches of documents for a single query.
encode(encoding: BatchEncoding, expansion: bool = False, pooling_strategy: 'first' | 'mean' | 'max' | 'sum' | None = None, mask_scoring_input_ids: Tensor | None = None) → BiEncoderEmbedding

Encodes a batched tokenized text sequences and returns the embeddings and scoring mask.

Parameters:

  • encoding (BatchEncoding) – Tokenizer encodings for the text sequence

  • expansion (bool, optional) – Whether mask expansion was applied to the text sequence, defaults to False

  • pooling_strategy (Literal['first', 'mean', 'max', 'sum'] | None, optional) – Strategy to pool token embeddings into a single embedding. If None no pooling is applied, defaults to None

  • mask_scoring_input_ids (torch.Tensor | None, optional) – Which token_ids to mask out during scoring, defaults to None

Returns:

Embeddings and scoring mask

Return type:

BiEncoderEmbedding

encode_doc(encoding: BatchEncoding) → BiEncoderEmbedding

Encodes tokenized documents.

Parameters:

encoding (BatchEncoding) – Tokenizer encodings for the documents

Returns:

Query embeddings and scoring mask

Return type:

BiEncoderEmbedding

encode_query(encoding: BatchEncoding) → BiEncoderEmbedding

Encodes tokenized queries.

Parameters:

encoding (BatchEncoding) – Tokenizer encodings for the queries

Returns:

Query embeddings and scoring mask

Return type:

BiEncoderEmbedding

forward(query_encoding: BatchEncoding | None, doc_encoding: BatchEncoding | None, num_docs: Sequence[int] | int | None = None) → BiEncoderOutput

Embeds queries and/or documents and computes relevance scores between them if both are provided.

Parameters:

  • query_encoding (BatchEncoding | None) – Tokenizer encodings for the queries

  • doc_encoding (BatchEncoding | None) – Tokenizer encodings for the documents

  • num_docs (Sequence[int] | int | None, optional) – Specifies how many documents are passed per query. If a sequence of integers, len(num_doc) should be equal to the number of queries and sum(num_docs) equal to the number of documents, i.e., the sequence contains one value per query specifying the number of documents for that query. If an integer, assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing the number of documents by the number of queries, defaults to None

Returns:

Output of the model

Return type:

BiEncoderOutput

get_output_embeddings() → Module | None

Returns the output embeddings of the model for tieing the input and output embeddings. Returns None if no MLM head is used for projection.

Returns:

Output embeddings of the model

Return type:

torch.nn.Module | None

score(query_embeddings: BiEncoderEmbedding, doc_embeddings: BiEncoderEmbedding, num_docs: Sequence[int] | int | None = None) → Tensor

Compute relevance scores between queries and documents.

Parameters:

  • query_embeddings (BiEncoderEmbedding) – Embeddings and scoring mask for the queries

  • doc_embeddings (BiEncoderEmbedding) – Embeddings and scoring mask for the documents

  • num_docs (Sequence[int] | int | None, optional) – Specifies how many documents are passed per query. If a sequence of integers, len(num_doc) should be equal to the number of queries and sum(num_docs) equal to the number of documents, i.e., the sequence contains one value per query specifying the number of documents for that query. If an integer, assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing the number of documents by the number of queries, defaults to None

Returns:

Relevance scores

Return type:

torch.Tensor

scoring_mask(encoding: BatchEncoding, expansion: bool = False, pooling_strategy: 'first' | 'mean' | 'max' | 'sum' | None = None, mask_scoring_input_ids: Tensor | None = None) → Tensor

Computes a scoring for batched tokenized text sequences which is used in the scoring function to mask out vectors during scoring.

Parameters:

  • encoding (BatchEncoding) – Tokenizer encodings for the text sequence

  • expansion (bool, optional) – Whether or not mask expansion was applied to the tokenized sequence, defaults to False

  • pooling_strategy (Literal['first', 'mean', 'max', 'sum'] | None, optional) – Which pooling strategy is pool the embeddings, defaults to None

  • mask_scoring_input_ids (torch.Tensor | None, optional) – Sequence of token_ids which should be masked during scoring, defaults to None

Returns:

Scoring mask

Return type:

torch.Tensor