File size: 1,846 Bytes
0cf964b | 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 | """Chronos-style tokenizer: mean-scale a window, bin it, reserve two special ids."""
import torch
PAD_ID = 0
EOS_ID = 1
SPECIAL_TOKENS = 2
class Tokenizer:
def __init__(self, n_bins: int= 4096, clip: float =15.0) -> None:
self.n_bins = n_bins
self.clip = clip
self.bin_width = 2.0 * clip / n_bins
@property
def vocab_size(self) -> int:
return self.n_bins + SPECIAL_TOKENS
def centres(self, device: torch.device | None = None) -> torch.Tensor:
half = self.bin_width / 2.0
return torch.linspace(-self.clip + half, self.clip - half, self.n_bins, device=device)
def mean_scale(self, x: torch.Tensor) -> torch.Tensor:
# Gaps are NaN and must not poison the mean; a dead window scales by 1 so z stays finite.
scale = torch.nanmean(x.abs(), dim=-1, keepdim=True)
return torch.where(scale > 0.0, scale, torch.ones_like(scale))
def encode(
self, x: torch.Tensor, scale: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
if scale is None:
scale = self.mean_scale(x)
gap = torch.isnan(x)
z = torch.nan_to_num(x / scale, nan=0.0).clamp(-self.clip, self.clip)
# z == clip lands on bin n_bins, one past the end, so the top bin absorbs it.
bins = torch.floor((z + self.clip) / self.bin_width).long().clamp(max=self.n_bins - 1)
ids = torch.where(gap, torch.full_like(bins, PAD_ID), bins + SPECIAL_TOKENS)
return ids, scale
def decode(self, ids: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
centres = self.centres(device=ids.device)
values = centres[(ids - SPECIAL_TOKENS).clamp(min=0)]* scale
return torch.where(ids < SPECIAL_TOKENS, torch.full_like(values, float("nan")), values) |