Testing utils

compute_biobert_loss_singlegpu(trainer, pl_module)

Computes the loss for BioBert models on a single GPU.

This will not function in multi-gpu settings nor with models that do not conform to BioBert.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The Lightning Trainer object.	required
pl_module	LightningModule	The LightningModule being trained.	required

Returns:

Name	Type	Description
float		The mean loss.

See Also: - :class: BioBertModel

Source code in bionemo/llm/model/biobert/testing_utils.py

def compute_biobert_loss_singlegpu(trainer: pl.Trainer, pl_module: pl.LightningModule):
    """Computes the loss for BioBert models on a single GPU.

    This will not function in multi-gpu settings nor with models that do not conform to BioBert.

    Args:
        trainer (pl.Trainer): The Lightning Trainer object.
        pl_module (pl.LightningModule): The LightningModule being trained.

    Returns:
        float: The mean loss.

    See Also:
    - :class: BioBertModel
    """
    model = pl_module
    dl = trainer.datamodule.val_dataloader()

    n, loss = -1, 0.0
    model.eval()
    # batch = next(iter(dl))
    batch = model.data_step(iter(dl))
    result = model(
        input_ids=batch["text"].cuda(),  # 'tokens' also a valid input for MockGPTDataModule
        attention_mask=batch["attention_mask"].cuda(),
    )
    loss_mask = batch["loss_mask"].cuda()
    # Not guaranteed i guess?
    logits = result["token_logits"]
    target = batch["labels"].cuda()
    loss += F.cross_entropy(logits[loss_mask].float(), target[loss_mask], reduction="sum")
    n += loss_mask.sum()

    mean_loss: float = (loss / n).detach().cpu().numpy().item()
    model.train()
    return mean_loss