ApproxNDCG

class lightning_ir.loss.approximate.ApproxNDCG(temperature: float = 1, scale_gains: bool = True)[source]

Bases: ApproxLossFunction

Approximate NDCG loss function for ranking tasks. Originally proposed in: Cumulated Gain-Based Evaluation of IR Techniques

__init__(temperature: float = 1, scale_gains: bool = True)[source]

Initialize the ApproxNDCG loss function.

Parameters:

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

  • scale_gains (bool) – Whether to scale the gains. Defaults to True.

Methods

__init__([temperature, scale_gains])

Initialize the ApproxNDCG loss function.

compute_loss(output, batch)

Compute the ApproxNDCG loss.

get_dcg(ranks, targets[, k, scale_gains])

Compute the Discounted Cumulative Gain (DCG) for the given ranks and targets.

get_ndcg(ranks, targets[, k, scale_gains, ...])

Compute the Normalized Discounted Cumulative Gain (NDCG) for the given ranks and targets.
compute_loss(output: LightningIROutput, batch: TrainBatch) → torch.Tensor[source]

Compute the ApproxNDCG 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_dcg(ranks: Tensor, targets: Tensor, k: int | None = None, scale_gains: bool = True) → Tensor[source]

Compute the Discounted Cumulative Gain (DCG) 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 DCG for the top k items.

  • scale_gains (bool) – Whether to scale the gains. Defaults to True.

Returns:

The computed DCG values.

Return type:

torch.Tensor

static get_ndcg(ranks: Tensor, targets: Tensor, k: int | None = None, scale_gains: bool = True, optimal_targets: Tensor | None = None) → Tensor[source]

Compute the Normalized Discounted Cumulative Gain (NDCG) 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) – Cutoff for the ranks. If provided, only computes NDCG for the top k items. Defaults to None.

  • scale_gains (bool) – Whether to scale the gains. Defaults to True.

  • optimal_targets (torch.Tensor | None) – Optional tensor of optimal targets for normalization. If None, uses the targets. Defaults to None.

Returns:

The computed NDCG 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