Spaces:
Running on Zero
Running on Zero
File size: 3,508 Bytes
45e7dfb | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """
spike_tokenizer.py -- HuggingFace-compatible wrapper for the custom
byte-level "length-max" (greedy longest-match) tokenizer in tokenizer.json.
The raw tokenizer.json is NOT a HuggingFace `tokenizers` file; it is a plain
dict {vocab, vocab_size, max_token_len, algorithm:"length-max"}. This wrapper
makes it loadable by AutoTokenizer.from_pretrained / save_pretrained and
exposes encode/decode + the bos/eos/pad/unk ids the training scripts expect.
Encoding scheme (verified): byte-level. Text is UTF-8 encoded, each byte mapped
to its latin-1 character, then greedily matched against the vocab using the
longest key that matches at each position (max key length = max_token_len).
"""
import json, os
from typing import List, Optional
from transformers import PreTrainedTokenizer
class SpikeTokenizer(PreTrainedTokenizer):
vocab_files_names = {"vocab_file": "tokenizer.json"}
model_input_names = ["input_ids"]
def __init__(self, vocab_file=None, **kwargs):
with open(vocab_file, "r", encoding="utf-8") as f:
data = json.load(f)
self._vocab = data["vocab"] # str -> id
self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
self.max_token_len = int(data.get("max_token_len", 24))
# length-bucketed keys for fast greedy match (longest length first)
self._lengths = sorted({len(k) for k in self._vocab}, reverse=True)
kwargs.setdefault("bos_token", "<bos>")
kwargs.setdefault("eos_token", "<eos>")
kwargs.setdefault("unk_token", "<unk>")
kwargs.setdefault("pad_token", "<pad>")
super().__init__(**kwargs)
@property
def vocab_size(self) -> int:
return len(self._vocab)
def get_vocab(self):
return dict(self._vocab)
# --- core byte-level greedy tokenization ---
def _tokenize(self, text: str) -> List[str]:
s = text.encode("utf-8").decode("latin-1") # one char per byte
out, i, n = [], 0, len(s)
while i < n:
matched = None
hi = min(self.max_token_len, n - i)
for L in range(hi, 0, -1):
sub = s[i:i + L]
if sub in self._vocab:
matched = sub
break
if matched is None: # single byte always exists in vocab
matched = s[i]
out.append(matched)
i += len(matched)
return out
def _convert_token_to_id(self, token: str) -> int:
return self._vocab.get(token, self._vocab["<unk>"])
def _convert_id_to_token(self, index: int) -> str:
return self._ids_to_tokens.get(index, "<unk>")
def convert_tokens_to_string(self, tokens: List[str]) -> str:
specials = {"<pad>", "<unk>", "<bos>", "<eos>"}
byte_str = "".join(t for t in tokens if t not in specials)
return byte_str.encode("latin-1").decode("utf-8", errors="replace")
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
os.makedirs(save_directory, exist_ok=True)
fn = (filename_prefix + "-" if filename_prefix else "") + "tokenizer.json"
path = os.path.join(save_directory, fn)
with open(path, "w", encoding="utf-8") as f:
json.dump({"vocab": self._vocab, "vocab_size": self.vocab_size,
"max_token_len": self.max_token_len,
"algorithm": "length-max"}, f, ensure_ascii=False)
return (path,)
|