sampling.py

import torch


def rep_penalty(logits: torch.Tensor, history_ids: torch.Tensor, penalty: float) -> torch.Tensor:
    # HuggingFace convention: penalty>1 discourages, penalty<1 encourages.
    # For each id in history: logit /= penalty if logit>0 else logit *= penalty.
    # logits: [B, V], history_ids: [B, T] (may be empty).
    if penalty == 1.0 or history_ids.numel() == 0:
        return logits
    B, V = logits.shape
    gathered = logits.gather(1, history_ids)  # [B, T]
    scaled = torch.where(gathered > 0, gathered / penalty, gathered * penalty)
    logits = logits.scatter(1, history_ids, scaled)
    return logits


def sample_logits(logits_1d: torch.Tensor,
                  temperature: float,
                  top_p: float,
                  min_p: float) -> torch.Tensor:
    # logits_1d: [B, V]. Applies (in order): temperature, min-p floor,
    # nucleus top-p, then multinomial sample. Returns [B] long.
    if temperature != 1.0:
        logits_1d = logits_1d / max(temperature, 1e-6)
    probs = torch.softmax(logits_1d, dim=-1)

    if min_p > 0.0:
        # Keep only tokens with prob >= min_p * max_prob (row-wise).
        top = probs.max(dim=-1, keepdim=True).values
        keep = probs >= (min_p * top)
        probs = probs * keep.to(probs.dtype)
        probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(1e-12)

    if 0.0 < top_p < 1.0:
        sorted_p, sorted_idx = torch.sort(probs, dim=-1, descending=True)
        cum = torch.cumsum(sorted_p, dim=-1)
        # Keep smallest set whose cumulative mass >= top_p. Shift by one so
        # the boundary token itself is kept.
        keep_sorted = cum - sorted_p <= top_p
        sorted_p = sorted_p * keep_sorted.to(sorted_p.dtype)
        sorted_p = sorted_p / sorted_p.sum(dim=-1, keepdim=True).clamp_min(1e-12)
        probs = torch.zeros_like(probs).scatter(1, sorted_idx, sorted_p)

    return torch.multinomial(probs, num_samples=1).squeeze(-1)