File size: 2,261 Bytes
5e41831 | 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 | import json
import os
import shutil
import torch
from transformers import PreTrainedTokenizerFast
class PrunedTokenizer(PreTrainedTokenizerFast):
def set_id_map(self, old_to_new: dict[int, int]) -> None:
size = max(int(k) for k in old_to_new) + 1
id_list = [0] * size
for old_id, new_id in old_to_new.items():
id_list[int(old_id)] = int(new_id)
self.init_kwargs["pruned_id_map"] = id_list
self._map_tensor = torch.tensor(id_list, dtype=torch.long)
def _ensure_map(self) -> None:
if getattr(self, "_map_tensor", None) is not None:
return
id_list = self.init_kwargs.get("pruned_id_map")
self._map_tensor = torch.tensor(id_list, dtype=torch.long) if id_list else None
def _encode_plus(self, *args, **kwargs):
encoding = super()._encode_plus(*args, **kwargs)
self._ensure_map()
if self._map_tensor is not None and "input_ids" in encoding:
ids = encoding["input_ids"]
if isinstance(ids, torch.Tensor):
remapped = self._map_tensor.to(ids.device)[ids.long()].to(ids.dtype)
else: # list / nested list / numpy array
remapped = self._map_tensor[torch.as_tensor(ids, dtype=torch.long)].tolist()
encoding["input_ids"] = remapped
return encoding
def save_pretrained(
self,
save_directory,
legacy_format=None,
filename_prefix=None,
push_to_hub=False,
**kwargs,
):
result = super().save_pretrained(
save_directory,
legacy_format=legacy_format,
filename_prefix=filename_prefix,
push_to_hub=push_to_hub,
**kwargs,
)
config_path = os.path.join(save_directory, "tokenizer_config.json")
with open(config_path) as f:
cfg = json.load(f)
cfg["auto_map"] = {"AutoTokenizer": [None, "pruning_tokenizer.PrunedTokenizer"]}
with open(config_path, "w") as f:
json.dump(cfg, f, indent=2)
# ship this file alongside the model so trust_remote_code=True finds the class
shutil.copy(__file__, os.path.join(save_directory, "pruning_tokenizer.py"))
return result
|