File size: 9,133 Bytes
e0ceb5c | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | """
Python wrapper around the C++ tokenizer with Hugging Face PreTrainedTokenizer compatibility.
"""
import json
import os
import shutil
from pathlib import Path
from typing import List, Union, Optional, Dict, Any
try:
from ._cpp_tokenizer import Tokenizer as _CTokenizer
except ImportError:
raise ImportError(
"C++ extension not built. Run: "
"pip install pybind11 && python setup.py build_ext --inplace"
)
class CurriculumTokenizer:
"""
Fast BPE tokenizer with curriculum training support.
Compatible with Hugging Face Transformers tokenization interface.
"""
def __init__(self, vocab_file: str, merges_file: str, **kwargs):
"""
Args:
vocab_file: Path to vocab.json
merges_file: Path to merges.txt
**kwargs: Additional config (special_tokens, etc.)
"""
self._tok = _CTokenizer()
vocab_path = os.path.abspath(vocab_file)
merges_path = os.path.abspath(merges_file)
self._vocab_path = vocab_path
self._merges_path = merges_path
self._config = dict(kwargs)
if not os.path.exists(vocab_path):
raise FileNotFoundError(f"Vocab file not found: {vocab_path}")
if not os.path.exists(merges_path):
raise FileNotFoundError(f"Merges file not found: {merges_path}")
ok = self._tok.load(vocab_path, merges_path)
if not ok:
raise RuntimeError("Failed to load tokenizer. Check vocab.json and merges.txt format.")
native_vocab_size = self._tok.vocab_size()
config_vocab_size = kwargs.get("vocab_size")
self.special_tokens = list(kwargs.get("special_tokens", ["<|endoftext|>"]))
if len(set(self.special_tokens)) != len(self.special_tokens):
raise ValueError("Duplicate special-token names")
declared_ids = kwargs.get("special_token_ids", kwargs.get("reserved_token_ids", {}))
if set(declared_ids) - set(self.special_tokens):
raise ValueError("Configured control IDs contain an undeclared special token")
self._special_token_ids = {
token: declared_ids.get(token, self._tok.token_to_id(token))
for token in self.special_tokens
}
self._special_aliases = {}
for token, token_id in kwargs.get("compatibility_reserved_ids", {}).items():
if token_id in self._special_aliases:
raise ValueError("Duplicate decode-only alias ID")
self._special_aliases[token_id] = token
inferred_size = max([native_vocab_size, *[i + 1 for i in self._special_token_ids.values()], *[i + 1 for i in self._special_aliases]])
expected_size = int(config_vocab_size) if config_vocab_size is not None else inferred_size
self._tok.configure_special_tokens(self._special_token_ids, self._special_aliases, expected_size)
self.vocab_size = self._tok.vocab_size()
self._skip_special_ids = set(self._special_token_ids.values()) | set(self._special_aliases)
# Build vocab mapping for Python access
self._vocab = {}
self._id_to_token = {}
for i in range(self.vocab_size):
token = self._tok.id_to_token(i)
if token:
if i not in self._special_aliases:
self._vocab[token] = i
self._id_to_token[i] = token
# Special token IDs
self.eos_token = kwargs.get("eos_token", "<|endoftext|>")
eos_id = self._tok.token_to_id(self.eos_token) if self.eos_token is not None else -1
self.eos_token_id = eos_id if eos_id >= 0 else None
self.pad_token = kwargs.get("pad_token", "<|endoftext|>")
pad_id = self._tok.token_to_id(self.pad_token) if self.pad_token is not None else -1
self.pad_token_id = pad_id if pad_id >= 0 else None
@classmethod
def from_pretrained(cls, model_dir: str, **kwargs):
"""Load tokenizer from a directory containing vocab.json and merges.txt."""
model_dir = Path(model_dir)
vocab_file = model_dir / "vocab.json"
merges_file = model_dir / "merges.txt"
# Load config if present
config_file = model_dir / "tokenizer_config.json"
if config_file.exists():
with open(config_file, encoding="utf-8") as f:
config = json.load(f)
kwargs = {**config, **kwargs}
return cls(str(vocab_file), str(merges_file), **kwargs)
def encode(
self,
text: Union[str, List[str]],
add_special_tokens: bool = True,
**kwargs
) -> Union[List[int], Dict[str, List[int]]]:
"""
Encode text to token IDs.
Args:
text: String or list of strings to encode
add_special_tokens: Whether to add EOS token
Returns:
List of token IDs, or dict with "input_ids" key (HF style)
"""
if isinstance(text, list):
return self.encode_batch(text, add_special_tokens=add_special_tokens, **kwargs)
ids = self._tok.encode(text)
if add_special_tokens and self.eos_token_id is not None:
ids.append(self.eos_token_id)
return ids
def encode_batch(
self,
texts: List[str],
add_special_tokens: bool = True,
num_threads: int = 0,
**kwargs
) -> Dict[str, List[List[int]]]:
"""Encode in input order; 0 selects up to four native workers for large batches."""
results = self._tok.encode_batch(texts, num_threads=num_threads)
if add_special_tokens and self.eos_token_id is not None:
for ids in results:
ids.append(self.eos_token_id)
return {"input_ids": results}
def set_cache_capacity(self, capacity: int) -> None:
"""Bound cached exact pieces; zero disables caching."""
self._tok.set_cache_capacity(capacity)
def clear_cache(self) -> None:
"""Clear cached encodings and reset hit/miss counters."""
self._tok.clear_cache()
def cache_info(self) -> Dict[str, int]:
"""Return piece-cache statistics; payload bytes exclude container overhead."""
return self._tok.cache_info()
def decode(
self,
ids: Union[List[int], List[List[int]]],
skip_special_tokens: bool = True,
**kwargs
) -> Union[str, List[str]]:
"""
Decode token IDs to text.
Args:
ids: List of token IDs or batch of lists
skip_special_tokens: Whether to remove special tokens
Returns:
Decoded string or list of strings
"""
if ids and isinstance(ids[0], list):
return [self.decode_single(seq, skip_special_tokens) for seq in ids]
return self.decode_single(ids, skip_special_tokens)
def decode_single(self, ids: List[int], skip_special_tokens: bool) -> str:
"""Decode a single sequence."""
if skip_special_tokens and self._skip_special_ids:
ids = [id for id in ids if id not in self._skip_special_ids]
return self._tok.decode(ids)
def __call__(self, text: Union[str, List[str]], **kwargs):
"""Hugging Face style __call__ method."""
result = self.encode(text, **kwargs)
if isinstance(result, dict):
return result
return {"input_ids": result}
def convert_tokens_to_ids(self, tokens: Union[str, List[str]]) -> Union[int, List[int]]:
"""Convert token strings to IDs."""
fallback = self.eos_token_id if self.eos_token_id is not None else 0
if isinstance(tokens, str):
native_id = self._tok.token_to_id(tokens)
if native_id >= 0:
return native_id
return self._vocab.get(tokens, fallback)
return [self.convert_tokens_to_ids(token) for token in tokens]
def convert_ids_to_tokens(self, ids: Union[int, List[int]]) -> Union[str, List[str]]:
"""Convert IDs to token strings."""
if isinstance(ids, int):
return self._id_to_token.get(ids, "")
return [self._id_to_token.get(i, "") for i in ids]
def save_pretrained(self, save_dir: str):
"""Save tokenizer files to a directory."""
os.makedirs(save_dir, exist_ok=True)
config = {**self._config,
"vocab_size": self.vocab_size,
"special_tokens": self.special_tokens,
"special_token_ids": self._special_token_ids,
"eos_token": self.eos_token,
"pad_token": self.pad_token,
}
with open(os.path.join(save_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
json.dump(config, f, indent=2)
for source, name in [(self._vocab_path, "vocab.json"), (self._merges_path, "merges.txt")]:
destination = os.path.abspath(os.path.join(save_dir, name))
if source != destination:
shutil.copyfile(source, destination)
def __len__(self):
return self.vocab_size
def __repr__(self):
return f"CurriculumTokenizer(vocab_size={self.vocab_size})"
|