File size: 1,602 Bytes
44ed717 | 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 | """Pruned-vocab id remap for ALF-emb-micro."""
from __future__ import annotations
from pathlib import Path
class RemapTokenizer:
def __init__(self, tok_dir: str | Path, keep_ids: list[int], unk_new: int | None = None):
from transformers import AutoTokenizer
import numpy as np
import torch
self.base = AutoTokenizer.from_pretrained(str(tok_dir))
self.map = {old: new for new, old in enumerate(keep_ids)}
unk_old = self.base.unk_token_id
self.unk_new = unk_new if unk_new is not None else self.map.get(unk_old, 0)
vocab = int(getattr(self.base, "vocab_size", 0) or 0)
size = max(vocab, max(keep_ids) + 1 if keep_ids else 1)
lut = np.full(size, self.unk_new, dtype=np.int64)
for old, new in self.map.items():
if 0 <= old < size:
lut[old] = new
self._lut_np = lut
self._lut = torch.from_numpy(lut)
def __call__(self, texts, **kwargs):
enc = self.base(texts, **kwargs)
ids = enc["input_ids"]
if hasattr(ids, "clamp"):
lut = self._lut.to(device=ids.device)
safe = ids.clamp(0, lut.numel() - 1).long()
enc["input_ids"] = lut[safe].to(dtype=ids.dtype)
else:
import numpy as np
arr = np.asarray(ids, dtype=np.int64)
mapped = self._lut_np[np.clip(arr, 0, len(self._lut_np) - 1)]
tensors = kwargs.get("return_tensors")
enc["input_ids"] = mapped if tensors in ("np", "pt") or hasattr(ids, "shape") else mapped.tolist()
return enc
|