ScoreBasedInBatchLossFunction

class lightning_ir.loss.in_batch.ScoreBasedInBatchLossFunction(min_target_diff: float, max_num_neg_samples: int | None = None)[source]

Bases: InBatchLossFunction

Base class for in-batch loss functions that compute in-batch indices based on scores.

__init__(min_target_diff: float, max_num_neg_samples: int | None = None)[source]

Initialize the ScoreBasedInBatchLossFunction.

Parameters:

  • min_target_diff (float) – Minimum target difference for negative sampling.

  • max_num_neg_samples (int | None) – Maximum number of negative samples.

Methods

__init__(min_target_diff[, max_num_neg_samples])

Initialize the ScoreBasedInBatchLossFunction.
abstract compute_loss(output: LightningIROutput, *args, **kwargs) → torch.Tensor

Compute the loss for the given output.

Parameters:

output (LightningIROutput) – The output from the model.

Returns:

The computed loss.

Return type:

torch.Tensor

get_ib_idcs(output: LightningIROutput, batch: TrainBatch) → Tuple[torch.Tensor, torch.Tensor]

Get in-batch indices for positive and negative samples.

Parameters:

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

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

Returns:

Indices of positive and negative samples.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Raises:

ValueError – If scores are not present in the output.

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