AstrAI/astrai/inference/runtime/sample.py

"""Composable sampling strategies for logit transformation.

Implements the Strategy pattern: each sampling technique
(temperature, top-k, top-p, frequency penalty) is a pluggable
strategy that can be composed into a pipeline.

All strategies accept both scalar and per-sample tensor
parameters, so a single pipeline works for any batch size.
"""

from abc import ABC, abstractmethod
from typing import List, Optional, Union

import torch
from torch import Tensor


class BaseSamplingStrategy(ABC):
    """Abstract base for a logit transformation strategy."""

    @abstractmethod
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Applies the strategy to logits.

        Args:
            logits: Raw logits tensor (batch, vocab_size).
            filter_value: Value assigned to filtered-out positions.
            input_ids: Previously generated token IDs ``[batch, seq_len]``,
                padded with 0. Used by frequency penalty.
            input_mask: Boolean mask ``[batch, seq_len]``, True for real
                tokens, False for padding. Used to exclude padding from
                penalty computation.

        Returns:
            Transformed logits tensor.
        """
        raise NotImplementedError

    @property
    def preserves_argmax(self) -> bool:
        """Whether ``apply`` never moves the argmax token.

        Conservative default: strategies must opt in. The greedy
        short-circuit in :class:`SamplingPipeline` asks this
        polymorphically, so a new strategy that can move the argmax
        automatically disables it — no isinstance bookkeeping.
        """
        return False

    @property
    def is_greedy(self) -> bool:
        """Whether this strategy collapses sampling onto the argmax token."""
        return False


class TemperatureStrategy(BaseSamplingStrategy):
    """Divides logits by temperature to control randomness.

    Args:
        temperature: Scalar or ``[batch]`` tensor.
    """

    def __init__(self, temperature: Union[float, Tensor] = 1.0):
        self.temperature = temperature

    @staticmethod
    def is_greedy_temperature(temperature: Union[float, Tensor]) -> bool:
        if isinstance(temperature, Tensor):
            return bool((temperature == 0).all())
        return temperature == 0

    @property
    def is_greedy(self) -> bool:
        return self.is_greedy_temperature(self.temperature)

    @property
    def preserves_argmax(self) -> bool:
        # Scaling by a positive constant (1/t, clamped away from zero)
        # preserves logit order; t=0 degenerates onto the argmax itself.
        return True

    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        t = self.temperature
        if isinstance(t, Tensor):
            t = t.to(logits.device, non_blocking=True).view(-1, 1)
            t = torch.clamp(t, min=1e-8)
            if (t != 1.0).any():
                logits = logits / t
        elif t != 1.0:
            logits = logits / max(t, 1e-8)
        return logits


class TopKStrategy(BaseSamplingStrategy):
    """Keeps only the top-k logits, setting the rest to filter_value.

    Args:
        top_k: Scalar or ``[batch]`` tensor (0 disables).
    """

    @property
    def preserves_argmax(self) -> bool:
        # The argmax token always ranks first, so any k >= 1 keeps it.
        return True

    def __init__(self, top_k: Union[int, Tensor] = 0):
        self.top_k = top_k

    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        tk = self.top_k
        if isinstance(tk, Tensor):
            tk = tk.to(logits.device, non_blocking=True).long().clamp(min=0)
            max_k = int(tk.max().item())
            if max_k <= 0:
                return logits
            max_k = min(max_k, logits.size(-1))
            values, _ = torch.topk(logits, max_k, dim=-1)
            per_row_k = tk.clamp(max=max_k)
            thresholds = torch.full_like(logits[..., -1:], -float("inf"))
            positive = per_row_k > 0
            if positive.any():
                row_idx = torch.arange(logits.size(0), device=logits.device)[positive]
                thresholds[positive] = values[
                    row_idx, per_row_k[positive] - 1
                ].unsqueeze(-1)
            logits[logits < thresholds] = filter_value
            return logits
        if tk > 0:
            k = min(tk, logits.size(-1))
            thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
            logits[logits < thresholds] = filter_value
        return logits


class TopPStrategy(BaseSamplingStrategy):
    """Nucleus (top-p) filtering: keeps the smallest set of tokens whose
    cumulative probability exceeds top_p.

    Args:
        top_p: Scalar or ``[batch]`` tensor (1.0 disables).
    """

    @property
    def preserves_argmax(self) -> bool:
        # Nucleus filtering always keeps the highest-probability token.
        return True

    def __init__(self, top_p: Union[float, Tensor] = 1.0):
        self.top_p = top_p

    def _apply(
        self, logits: Tensor, top_p: Union[float, Tensor], filter_value: float
    ) -> Tensor:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        remove = cum_probs > top_p
        remove[..., 1:] = remove[..., :-1].clone()
        remove[..., 0] = False
        mask = torch.zeros_like(logits, dtype=torch.bool)
        mask.scatter_(1, sorted_indices, remove)
        logits[mask] = filter_value
        return logits

    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        tp = self.top_p
        if isinstance(tp, Tensor):
            tp = tp.to(logits.device, non_blocking=True)
            if (tp < 1.0).any():
                logits = self._apply(logits, tp.view(-1, 1), filter_value)
        elif tp < 1.0:
            logits = self._apply(logits, tp, filter_value)
        return logits


class FrequencyPenaltyStrategy(BaseSamplingStrategy):
    """Penalizes tokens based on how many times they appeared in history.

    Subtracts ``penalty * count(token)`` from each token's logit, where
    ``count(token)`` is the number of occurrences in the generation history
    (prompt + output). A penalty of ``0.0`` disables the strategy.

    Unlike repetition penalty (which only checks *presence*), frequency
    penalty scales linearly with occurrence count: the first use is
    penalized once, the third use three times. This allows natural
    repetition of common words while suppressing degenerate loops.

    Reference: OpenAI API ``frequency_penalty`` parameter.

    Args:
        penalty: Scalar or ``[batch]`` tensor (0.0 disables, range -2.0~2.0).
    """

    def __init__(self, penalty: Union[float, Tensor] = 0.0):
        self.penalty = penalty

    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        if input_ids is None:
            return logits

        p = self.penalty
        if isinstance(p, Tensor):
            p = p.to(logits.device, non_blocking=True).view(-1)
            if (p == 0.0).all():
                return logits
        elif p == 0.0:
            return logits

        input_ids = input_ids.to(logits.device, non_blocking=True)
        if input_mask is not None:
            input_mask = input_mask.to(logits.device, non_blocking=True)

        batch_sz = input_ids.shape[0]
        vocab_size = logits.size(-1)

        # Sync-free update: map each history token to a flat
        # ``row * vocab + token`` bucket (padding to one trailing sentinel
        # bucket), count with ``index_add_``, and subtract in one
        # elementwise pass. No nonzero/unique/boolean-mask indexing, so the
        # hot path never forces a device-host synchronization.
        row_offsets = (
            torch.arange(batch_sz, device=logits.device, dtype=torch.long).unsqueeze(1)
            * vocab_size
        )
        if input_mask is not None:
            flat = torch.where(
                input_mask,
                row_offsets + input_ids,
                torch.full_like(input_ids, batch_sz * vocab_size),
            )
        else:
            flat = row_offsets + input_ids
        flat = flat.reshape(-1)
        counts = torch.zeros(
            batch_sz * vocab_size + 1, device=logits.device, dtype=torch.float32
        )
        counts.index_add_(0, flat, torch.ones_like(flat, dtype=torch.float32))
        counts = counts[: batch_sz * vocab_size].view(batch_sz, vocab_size)
        if isinstance(p, Tensor):
            deltas = counts * p.to(torch.float32).view(-1, 1)
        else:
            deltas = counts * float(p)
        return logits - deltas.to(logits.dtype)


class SamplingPipeline(BaseSamplingStrategy):
    """Composes multiple sampling strategies into a single transformation.

    Strategies are applied sequentially in the order they are provided,
    matching the original temperature -> top-k -> top-p ordering.

    Usage::

        pipeline = SamplingPipeline([
            TemperatureStrategy(0.8),
            TopKStrategy(50),
            TopPStrategy(0.95),
        ])
        logits = pipeline.apply(logits)
        token = pipeline.sample(logits)       # softmax + multinomial
    """

    def __init__(self, strategies: List[BaseSamplingStrategy]):
        self.strategies = strategies

    @property
    def preserves_argmax(self) -> bool:
        # A composite preserves the argmax iff every stage does.
        return all(s.preserves_argmax for s in self.strategies)

    @property
    def is_greedy(self) -> bool:
        """Whether sampling always yields the argmax of the raw logits.

        True iff some stage forces greedy and no stage can move the
        argmax before or after it. Both facts are declared
        polymorphically by each strategy, so composing in a new strategy
        type (or a nested pipeline) updates this automatically.
        """
        return any(s.is_greedy for s in self.strategies) and self.preserves_argmax

    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        for strategy in self.strategies:
            logits = strategy.apply(logits, filter_value, input_ids, input_mask)
        return logits

    @torch.inference_mode()
    def sample(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
        return_logprobs: bool = False,
    ):
        """Apply strategies then sample (softmax + multinomial).

        Short-circuits to ``argmax`` when temperature is exactly 0
        (deterministic / greedy decode).

        Args:
            logits: Raw logits ``[batch, vocab_size]``.
            input_ids: Previously generated token IDs ``[batch, seq_len]``.
            input_mask: Boolean mask for ``input_ids`` padding.
            return_logprobs: If ``True``, return ``(tokens, logprobs)``
                where ``logprobs[i]`` is the log-probability of
                ``tokens[i]`` under the raw (pre-strategy) model
                distribution, matching training-side policy logprobs.

        Returns:
            Sampled token IDs ``[batch]``, or — when ``return_logprobs``
            is ``True`` — a ``(token_ids, chosen_logprobs)`` tuple.
        """
        if self.is_greedy:
            tokens = logits.argmax(dim=-1)
            if not return_logprobs:
                return tokens
            log_probs = torch.log_softmax(logits.float(), dim=-1)
            chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
            return tokens, chosen

        # Capture the raw distribution before the strategy pipeline runs:
        # top-k/top-p mutate the logits tensor in place, so computing this
        # after ``apply`` would read the filtered distribution instead of
        # the raw model distribution the caller documented.
        if return_logprobs:
            raw_log_probs = torch.log_softmax(logits.float(), dim=-1)

        transformed = self.apply(logits, filter_value, input_ids, input_mask)
        tokens = torch.multinomial(
            torch.softmax(transformed, dim=-1), num_samples=1
        ).squeeze(-1)
        if not return_logprobs:
            return tokens
        # Log-probabilities of the raw (pre-strategy) model distribution,
        # matching the training-side policy logprobs exactly: the behaviour
        # logprobs recorded for online RL must live in the same
        # distribution the trainer differentiates, not the
        # temperature/top-p filtered one tokens were drawn from.
        chosen = torch.gather(raw_log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
        return tokens, chosen


@torch.inference_mode()
def sample(
    logits: Tensor,
    temperature: Union[float, Tensor] = 1.0,
    top_k: Union[int, Tensor] = 0,
    top_p: Union[float, Tensor] = 1.0,
    frequency_penalty: Union[float, Tensor] = 0.0,
    input_ids: Optional[Tensor] = None,
    input_mask: Optional[Tensor] = None,
    filter_value: float = -float("inf"),
    return_logprobs: bool = False,
):
    """Apply sampling strategies then sample (softmax + multinomial).

    Shortcut for ``SamplingPipeline(...).sample(logits, return_logprobs=)``.

    When **temperature** is exactly 0 (scalar or single-element tensor)
    the function short-circuits to ``argmax`` for deterministic decode.

    When **frequency_penalty** is 0 (the common decode case), the entire
    frequency penalty computation — including the O(batch * vocab) count
    tensor allocation — is skipped.

    Args:
        logits: Raw logits ``[batch, vocab_size]``.
        frequency_penalty: Penalty per occurrence for repeated tokens
            (0.0 disables, range -2.0~2.0).
        input_ids: Previously generated token IDs ``[batch, seq_len]``.
        input_mask: Boolean mask for ``input_ids`` padding.
        return_logprobs: If ``True``, also return the log-probability
            of each sampled token under the raw (pre-strategy) model
            distribution — usable directly for RL rollout (PPO/GRPO
            importance ratios against the training-side policy logprobs).

    Returns:
        Sampled token IDs ``[batch]``, or — when ``return_logprobs`` is
        ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where
        ``chosen_logprobs`` has shape ``[batch]``.
    """
    has_freq = (
        (isinstance(frequency_penalty, Tensor) and (frequency_penalty != 0).any())
        if isinstance(frequency_penalty, Tensor)
        else frequency_penalty != 0
    )

    strategies: List[BaseSamplingStrategy] = []
    if has_freq:
        # Penalty first, on the raw logits (OpenAI semantics): applying it
        # after a temperature scaling would shrink it by the temperature
        # and annihilate it entirely at temperature=0.
        strategies.append(FrequencyPenaltyStrategy(frequency_penalty))
    strategies.extend(
        [
            TemperatureStrategy(temperature),
            TopKStrategy(top_k),
            TopPStrategy(top_p),
        ]
    )

    return SamplingPipeline(strategies).sample(
        logits,
        filter_value=filter_value,
        input_ids=input_ids,
        input_mask=input_mask,
        return_logprobs=return_logprobs,
    )