nemo_automodel.recipes.retrieval.train_cross_encoder

Module Contents

Classes

Name	Description
`TrainCrossEncoderRecipe`	-

Functions

Name	Description
`_validate_temperature_sources`	Reject simultaneous recipe-level and model-level temperature scaling.
`accuracy`	Returns (num_correct, batch_size) for top-1 accuracy.
`batch_mrr`	Returns sum of reciprocal ranks for the batch. Stays on GPU.
`main`	-

API

class nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe()

Bases: TrainBiEncoderRecipe

nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_batches,
    is_train: bool = True,
    modality_loss_buffers = None
)

Forward and backward pass for a single micro-batch.

nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._run_train_optim_step(
    batches,
    max_grad_norm = None
)

nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._run_validation_epoch(
    val_dataloader
)

Run validation for one epoch and compute loss, accuracy@1, and MRR.

nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe._validate_model(
    model: torch.nn.Module
) -> None

Validate the effective temperature applied by the constructed model.

nemo_automodel.recipes.retrieval.train_cross_encoder.TrainCrossEncoderRecipe.log_train_metrics(
    log_data: nemo_automodel.components.loggers.metric_logger.MetricsSample
)

nemo_automodel.recipes.retrieval.train_cross_encoder._validate_temperature_sources(
    recipe_temperature: float,
    model_temperature: float
) -> None

Reject simultaneous recipe-level and model-level temperature scaling.

Parameters:

recipe_temperature

float

Temperature applied by the cross-encoder recipe.

model_temperature

float

Temperature applied by the loaded model.

Raises:

ValueError: If both temperatures are non-unit.

nemo_automodel.recipes.retrieval.train_cross_encoder.accuracy(
    output: torch.Tensor,
    target: torch.Tensor
) -> tuple[torch.Tensor, int]

Returns (num_correct, batch_size) for top-1 accuracy.

nemo_automodel.recipes.retrieval.train_cross_encoder.batch_mrr(
    output: torch.Tensor,
    target: torch.Tensor
) -> torch.Tensor

Returns sum of reciprocal ranks for the batch. Stays on GPU.

nemo_automodel.recipes.retrieval.train_cross_encoder.main(
    default_config_path = 'examples/retrieval/cross_e...
)