SpongeBobFan2002/Zonosz

1import torch2 3 4def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):5    """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.6 7    Args:8        input (torch.Tensor): The input tensor containing probabilities.9        num_samples (int): Number of samples to draw.10        replacement (bool): Whether to draw with replacement or not.11    Keywords args:12        generator (torch.Generator): A pseudorandom number generator for sampling.13    Returns:14        torch.Tensor: Last dimension contains num_samples indices15            sampled from the multinomial probability distribution16            located in the last dimension of tensor input.17    """18 19    if num_samples == 1:20        q = torch.empty_like(input).exponential_(1, generator=generator)21        return torch.argmax(input / q, dim=-1, keepdim=True).to(torch.int64)22 23    input_ = input.reshape(-1, input.shape[-1])24    output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)25    output = output_.reshape(*list(input.shape[:-1]), -1)26    return output27 28 29def apply_top_k(30    probs: torch.Tensor,31    k: int,32) -> torch.Tensor:33    """Sample next token from top K values along the last dimension of the input probs tensor.34 35    Args:36        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.37        k (int): The k in “top-k”.38    Returns:39        torch.Tensor: Sampled tokens.40    """41    v, _ = torch.topk(probs, min(k, probs.size(-1)))42    pivot = v.select(-1, -1).unsqueeze(-1)43    probs = torch.where(probs < pivot, 0.0, probs)44    probs.div_(probs.sum(dim=-1, keepdim=True))45    return probs46 47 48def apply_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:49    """Sample next token from top P probabilities along the last dimension of the input probs tensor.50 51    Args:52        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.53        p (int): The p in “top-p”.54    Returns:55        torch.Tensor: Sampled tokens.56    """57    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)58    probs_sum = torch.cumsum(probs_sort, dim=-1)59    mask = probs_sum - probs_sort > p60    probs_sort *= (~mask).float()61    probs = probs.scatter(-1, probs_idx, probs_sort)62    probs.div_(probs.sum(dim=-1, keepdim=True))63    return probs64 65 66def apply_min_p(probs: torch.Tensor, min_p: float) -> torch.Tensor:67    """Sample next token using min-p sampling.68 69    Args:70        scores (torch.FloatTensor): Input logits with token candidates on the last dimension.71        min_p (float): Minimum token probability, scaled by the probability of the most likely token.72                       Must be between 0 and 1. Typical values are in the 0.01-0.2 range.73    Returns:74        torch.Tensor: Sampled tokens.75    """76    top_probs, _ = probs.max(dim=-1, keepdim=True)77    tokens_to_remove = probs < (min_p * top_probs)78    probs = probs.masked_fill(tokens_to_remove, 0.0)79    probs.div_(probs.sum(dim=-1, keepdim=True))80    return probs81 82 83def modify_logit_for_repetition_penalty(84    logits: torch.Tensor,85    generated_tokens: torch.Tensor,86    repetition_penalty: float,87    repetition_penalty_window: int,88):89    """See https://arxiv.org/abs/1909.0585890    Apply repetition penalty over a sliding window of the last `repetition_penalty_window` tokens.91    logits: (batch_size, n_codebooks, vocab_size)92    generated_tokens: (batch_size, n_codebooks, seq_len)93    """94    generated_tokens = generated_tokens[..., -repetition_penalty_window:]95    generated_tokens = generated_tokens.clamp_max(logits.shape[-1] - 1).to(torch.int64)96    rp = torch.full_like(logits, repetition_penalty)97    factors = torch.ones_like(logits).scatter_reduce(2, generated_tokens, rp, reduce="prod")98    return torch.where(logits <= 0, logits * factors, logits / factors)99 100 101def sample_from_logits(102    logits: torch.Tensor,103    temperature: float = 1.0,104    top_p: float = 0.0,105    top_k: int = 0,106    min_p: float = 0.0,107    generated_tokens: torch.Tensor | None = None,108    repetition_penalty: float = 3.0,109    repetition_penalty_window: float = 2,110) -> torch.Tensor:111    """Sample next token from logits using temperature, top-p, top-k, or min-p sampling.112 113    Args:114        logits (torch.Tensor): Input logits with token candidates on the last dimension.115        temperature (float): Sampling temperature. Lower temperature results in more deterministic samples.116        top_p (float): The p in “top-p”.117        top_k (int): The k in “top-k”.118        min_p (float): Minimum token probability, scaled by the probability of the most likely token.119                       Must be between 0 and 1. Typical values are in the 0.01-0.2 range.120 121    Returns:122        torch.Tensor: Sampled tokens.123    """124    if repetition_penalty != 1.0 and generated_tokens is not None:125        logits = modify_logit_for_repetition_penalty(logits, generated_tokens, repetition_penalty, repetition_penalty_window)126 127    if temperature > 0:128        probs = torch.softmax(logits / temperature, dim=-1)129 130        if top_p > 0:131            probs = apply_top_p(probs, top_p)132        if top_k > 0:133            probs = apply_top_k(probs, top_k)134        if min_p > 0:135            probs = apply_min_p(probs, min_p)136 137        next_token = multinomial(probs, num_samples=1)138    else:139        next_token = torch.argmax(logits, dim=-1, keepdim=True)140 141    return next_token  # [batch_size, num_codebooks, 1]142