RegularizationLossFunction

class lightning_ir.loss.base.RegularizationLossFunction(query_weight: float = 0.0001, doc_weight: float = 0.0001)[source]

Bases: EmbeddingLossFunction

Base class for regularization loss functions that operate on embeddings.

__init__(query_weight: float = 0.0001, doc_weight: float = 0.0001) → None[source]

Initialize the RegularizationLossFunction.

Parameters:

  • query_weight (float) – Weight for the query embeddings regularization. Defaults to 1e-4.

  • doc_weight (float) – Weight for the document embeddings regularization. Defaults to 1e-4.

Methods

__init__([query_weight, doc_weight])

Initialize the RegularizationLossFunction.

process_embeddings(output)

Process the embeddings from the output.
process_embeddings(output: BiEncoderOutput) → Tuple[torch.Tensor, torch.Tensor][source]

Process the embeddings from the output.

Parameters:

output (BiEncoderOutput) – The output from the model containing query and document embeddings.

Returns:

The processed query and document embeddings.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Raises:

  • ValueError – If query_embeddings are not present in the output.

  • ValueError – If doc_embeddings are not present in the output.