Function sequence_logprobs

src/post_training/rollout.py:268–284

Sequence-level summed log-prob over response tokens (used by DPO/KTO/ORPO). Returns ``(sum_logprob, n_tokens)`` each shape (B,). The per-token mean is ``sum_logprob / n_tokens.clamp(min=1)``.

(
    model,
    sequences: torch.Tensor,
    response_mask: torch.Tensor,
    *,
    temperature: float = 1.0,
    requires_grad: bool = True,
)

def sequence_logprobs(
    model,
    sequences: torch.Tensor,
    response_mask: torch.Tensor,
    *,
    temperature: float = 1.0,
    requires_grad: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Sequence-level summed log-prob over response tokens (used by DPO/KTO/ORPO).

    Returns ``(sum_logprob, n_tokens)`` each shape (B,). The per-token mean is
    ``sum_logprob / n_tokens.clamp(min=1)``.
    """
    lp, mask = compute_logprobs(model, sequences, response_mask, temperature=temperature, requires_grad=requires_grad)
    m = mask.to(lp.dtype)
    return (lp * m).sum(dim=-1), m.sum(dim=-1)


def sequence_entropy(model, sequences: torch.Tensor, temperature: float = 1.0) -> torch.Tensor: