Spaces:
Running
Running
File size: 3,380 Bytes
e319a7b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Sampling utilities for autoregressive generation."""
import torch
def length_to_mask(lengths: torch.Tensor, max_len: int | None = None) -> torch.Tensor:
"""Convert a tensor of sequence lengths to a boolean mask."""
assert len(lengths.shape) == 1
final_length = int(lengths.max().item()) if not max_len else max_len
final_length = max(final_length, 1)
return torch.arange(final_length, device=lengths.device)[None, :] < lengths[:, None]
def multinomial(
input: torch.Tensor, num_samples: int, replacement: bool = False, *, generator=None
) -> torch.Tensor:
"""torch.multinomial with arbitrary number of dimensions."""
input_ = input.reshape(-1, input.shape[-1])
output_ = torch.multinomial(
input_, num_samples=num_samples, replacement=replacement, generator=generator
)
output = output_.reshape(*list(input.shape[:-1]), -1)
return output
def sample_top_k(probs: torch.Tensor, k: int, num_samples: int = 1) -> torch.Tensor:
"""Sample from top-k probabilities."""
top_k_value, _ = torch.topk(probs, k, dim=-1)
min_value_top_k = top_k_value[..., [-1]]
probs = probs * (probs >= min_value_top_k).float()
probs = probs / probs.sum(dim=-1, keepdim=True)
return multinomial(probs, num_samples=num_samples)
def sample_top_p(probs: torch.Tensor, p: float, num_samples: int = 1) -> torch.Tensor:
"""Sample from nucleus (top-p) distribution."""
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
probs_sum = torch.cumsum(probs_sort, dim=-1)
mask = probs_sum - probs_sort > p
probs_sort = probs_sort * (~mask).float()
probs_sort = probs_sort / probs_sort.sum(dim=-1, keepdim=True)
next_token = multinomial(probs_sort, num_samples=num_samples)
return torch.gather(probs_idx, -1, next_token)
def sample_from_probs(
probs: torch.Tensor, top_p: float = 0.0, top_k: int = 0
) -> torch.Tensor:
"""Sample one token from probs, optionally filtered by top-p or top-k."""
if top_p > 0.0:
return sample_top_p(probs, top_p)
if top_k > 0:
return sample_top_k(probs, top_k)
return multinomial(probs, num_samples=1)
def sample_stratified(
probs: torch.Tensor,
special_token: int,
first_temp: float,
second_temp: float = 1.0,
top_p: float = 0.0,
top_k: int = 0,
) -> torch.Tensor:
"""Stratified sampling: first decide special vs. non-special, then sample among non-special."""
eps = 1e-12
probs_special = probs[..., special_token : special_token + 1].clamp(
min=eps, max=1 - eps
)
logits_two = torch.cat(
[torch.log(probs_special), torch.log(1 - probs_special)], dim=-1
)
logits_two = logits_two / max(first_temp, eps)
probs_two = torch.softmax(logits_two, dim=-1)
probs_special_temp = probs_two[..., 0:1]
next_token_is_special = torch.rand_like(probs_special_temp).lt(probs_special_temp)
denom = (1 - probs_special).clamp(min=eps)
new_probs = probs.clone() / denom
new_probs[..., special_token] = 0.0
if second_temp > 0:
log_new = torch.log(new_probs.clamp(min=eps)) / second_temp
new_probs = torch.softmax(log_new, dim=-1)
next_token = sample_from_probs(new_probs, top_p=top_p, top_k=top_k)
return torch.where(
next_token_is_special, torch.full_like(next_token, special_token), next_token
)
|