nemo_automodel.components.loss.infonce

Module Contents

Classes

Name	Description
`InfoNCEDistillLoss`	InfoNCE soft listwise distillation loss module.
`InfoNCELoss`	InfoNCE loss module with optional learnable temperature.

Functions

Name	Description
`infonce_distill_loss`	Soft listwise distillation on InfoNCE candidate sets.
`infonce_loss`	InfoNCE contrastive loss with optional hard negatives.

API

class nemo_automodel.components.loss.infonce.InfoNCEDistillLoss(
    temperature: float = 0.05,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True,
    divergence: str = 'kl',
    cross_device_negatives: bool = True
)

Bases: Module

InfoNCE soft listwise distillation loss module.

nemo_automodel.components.loss.infonce.InfoNCEDistillLoss.forward(
    student_queries: torch.Tensor,
    student_documents: torch.Tensor,
    teacher_queries: torch.Tensor,
    teacher_documents: torch.Tensor,
    student_hard_negatives: torch.Tensor | None = None,
    teacher_hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None
) -> torch.Tensor

class nemo_automodel.components.loss.infonce.InfoNCELoss(
    temperature: float = 0.05,
    learnable_temperature: bool = False,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True,
    cross_device_negatives: bool = True
)

Bases: Module

InfoNCE loss module with optional learnable temperature.

log_inv_tau

nemo_automodel.components.loss.infonce.InfoNCELoss.current_temperature() -> torch.Tensor

nemo_automodel.components.loss.infonce.InfoNCELoss.forward(
    queries: torch.Tensor,
    documents: torch.Tensor,
    hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None
) -> torch.Tensor

nemo_automodel.components.loss.infonce.infonce_distill_loss(
    student_queries: torch.Tensor,
    student_documents: torch.Tensor,
    teacher_queries: torch.Tensor,
    teacher_documents: torch.Tensor,
    student_hard_negatives: torch.Tensor | None = None,
    teacher_hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None,
    temperature: float | torch.Tensor = 0.05,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True,
    divergence: str = 'kl'
) -> torch.Tensor

Soft listwise distillation on InfoNCE candidate sets.

nemo_automodel.components.loss.infonce.infonce_loss(
    queries: torch.Tensor,
    documents: torch.Tensor,
    hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None,
    temperature: float | torch.Tensor = 0.05,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True
) -> torch.Tensor

InfoNCE contrastive loss with optional hard negatives.