File size: 1,688 Bytes
70cc4de | 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 | import json
import os
from typing import List, Optional, Tuple
from transformers import PreTrainedTokenizer
VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}
class TinyGPTTokenizer(PreTrainedTokenizer):
"""Character-level tokenizer: each of the 128 ASCII code points is its
own token, id == ord(char) — the exact scheme used by decode()/chr(t)
in the original training script."""
vocab_files_names = VOCAB_FILES_NAMES
model_input_names = ["input_ids", "attention_mask"]
def __init__(self, vocab_file: Optional[str] = None, **kwargs):
self._vocab = {chr(i): i for i in range(128)}
self._ids_to_tokens = {i: chr(i) for i in range(128)}
super().__init__(**kwargs)
@property
def vocab_size(self) -> int:
return len(self._vocab)
def get_vocab(self):
return dict(self._vocab)
def _tokenize(self, text: str, **kwargs) -> List[str]:
return list(text)
def _convert_token_to_id(self, token: str) -> int:
return self._vocab.get(token, self._vocab.get(" "))
def _convert_id_to_token(self, index: int) -> str:
return self._ids_to_tokens.get(index, " ")
def convert_tokens_to_string(self, tokens: List[str]) -> str:
return "".join(tokens)
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
filename = (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
vocab_path = os.path.join(save_directory, filename)
with open(vocab_path, "w", encoding="utf-8") as f:
json.dump(self._vocab, f, ensure_ascii=False, indent=2)
return (vocab_path,)
|