ApproxMRR

class lightning_ir.loss.approximate.ApproxMRR(temperature: float = 1)[source]

Bases: ApproxLossFunction

Approximate Mean Reciprocal Rank (MRR) loss function for ranking tasks.

__init__(temperature: float = 1)[source]

Initialize the ApproxMRR loss function.

Parameters:

temperature (float) – Temperature parameter for scaling the scores. Defaults to 1.

Methods

__init__([temperature])

Initialize the ApproxMRR loss function.

compute_loss(output, batch)

Compute the ApproxMRR loss.

get_mrr(ranks, targets[, k])

Compute the Mean Reciprocal Rank (MRR) for the given ranks and targets.
compute_loss(output: LightningIROutput, batch: TrainBatch) → torch.Tensor[source]

Compute the ApproxMRR 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

static get_approx_ranks(scores: Tensor, temperature: float) → Tensor

Compute approximate ranks from scores.

Parameters:

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

  • temperature (float) – Temperature parameter for scaling the scores.

Returns:

The computed approximate ranks.

Return type:

torch.Tensor

static get_mrr(ranks: Tensor, targets: Tensor, k: int | None = None) → Tensor[source]

Compute the Mean Reciprocal Rank (MRR) for the given ranks and targets.

Parameters:

  • ranks (torch.Tensor) – The ranks of the items.

  • targets (torch.Tensor) – The relevance scores of the items.

  • k (int | None) – Optional cutoff for the ranks. If provided, only computes MRR for the top k items.

Returns:

The computed MRR values.

Return type:

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