| """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:
|
|
|
| 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)
|
|
|
| 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) |
