ConstantMarginMSE

class lightning_ir.loss.pairwise.ConstantMarginMSE(margin: float = 1.0)[source]

Bases: MarginMSE

Constant Margin MSE loss for pairwise ranking tasks with a fixed margin.

__init__(margin: float = 1.0)[source]

Initialize the ConstantMarginMSE loss function.

Parameters:

margin (float) – The fixed margin value for the loss.

Methods

__init__([margin])

Initialize the ConstantMarginMSE loss function.
compute_loss(output: LightningIROutput, batch: TrainBatch) → torch.Tensor

Compute the MarginMSE loss.

Parameters:

  • output (LightningIROutput) – The output from the model containing scores.

  • batch (TrainBatch) – The training batch containing targets.

Returns:

The computed loss.

Return type:

torch.Tensor

Raises:

ValueError – If the margin type is invalid.

get_pairwise_idcs(targets: Tensor) → Tuple[Tensor, ...]

Get pairwise indices for positive and negative samples based on targets.

Parameters:

targets (torch.Tensor) – The targets tensor containing relevance labels.

Returns:

Indices of positive and negative samples.

Return type:

Tuple[torch.Tensor, …]

process_scores(output: LightningIROutput) → torch.Tensor

Process the scores from the output.

Parameters:

output (LightningIROutput) – The output from the model.

Returns:

The scores tensor.

Return type:

torch.Tensor

process_targets(scores: torch.Tensor, batch: TrainBatch) → torch.Tensor

Process the targets from the batch.

Parameters:

  • scores (torch.Tensor) – The scores tensor.

  • batch (TrainBatch) – The training batch.

Returns:

The processed targets tensor.

Return type:

torch.Tensor