InBatchLossFunction

class lightning_ir.loss.base.InBatchLossFunction(pos_sampling_technique: 'all' | 'first' = 'all', neg_sampling_technique: 'all' | 'first' | 'all_and_non_first' = 'all', max_num_neg_samples: int | None = None)[source]

Bases: LossFunction

Base class for in-batch loss functions that compute in-batch indices for positive and negative samples.

__init__(pos_sampling_technique: 'all' | 'first' = 'all', neg_sampling_technique: 'all' | 'first' | 'all_and_non_first' = 'all', max_num_neg_samples: int | None = None)[source]

Initialize the InBatchLossFunction.

Parameters:

  • pos_sampling_technique (Literal["all", "first"]) – Technique for positive sample sampling.

  • neg_sampling_technique (Literal["all", "first", "all_and_non_first"]) – Technique for negative sample sampling.

  • max_num_neg_samples (int | None) – Maximum number of negative samples to consider. If None, all negative samples are considered.

Raises:

ValueError – If the negative sampling technique is invalid for the given positive sampling technique.

Methods

__init__([pos_sampling_technique, ...])

Initialize the InBatchLossFunction.

get_ib_idcs(output, batch)

Get in-batch indices for positive and negative samples.
get_ib_idcs(output: LightningIROutput, batch: TrainBatch) → Tuple[torch.Tensor, torch.Tensor][source]

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.