Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / text / build_tokenizer.py: 0%
90 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Build the tokenizer used by the private LLM data stages."""
17import json
18from collections import OrderedDict
19from typing import Any
21from hyper_parallel.data.dataset_logging import get_dataset_logger
23logger = get_dataset_logger(__name__)
26class _DatasetTokenizer:
27 """Provide the stable tokenizer identity required by Dataset caches."""
29 def __init__(
30 self,
31 *tokenizer_paths: str,
32 vocab_size: int | None = None,
33 eod_token_id: int | None = None,
34 pad_token_id: int | None = None,
35 **tokenizer_options: Any,
36 ) -> None:
37 """Record the tokenizer identity and validate the special token configuration.
39 Args:
40 *tokenizer_paths: Tokenizer source paths recorded in the cache identity.
41 vocab_size: Pretokenized-corpus vocabulary size.
42 eod_token_id: Pretokenized-corpus end-of-document token ID.
43 pad_token_id: Pretokenized-corpus padding token ID, when one exists.
44 **tokenizer_options: Extra options recorded in the cache identity.
46 Raises:
47 ValueError: If the special token configuration is incomplete or out of range.
48 """
49 if (vocab_size is None) != (eod_token_id is None):
50 raise ValueError("vocab_size and eod_token_id must be provided together")
51 if vocab_size is not None and eod_token_id is not None:
52 if vocab_size <= 0:
53 raise ValueError("vocab_size must be positive")
54 if not 0 <= eod_token_id < vocab_size:
55 raise ValueError("eod_token_id must be within the tokenizer vocabulary")
56 if pad_token_id is not None and not 0 <= pad_token_id < vocab_size:
57 raise ValueError("pad_token_id must be within the tokenizer vocabulary")
58 tokenizer_options["vocab_size"] = vocab_size
59 tokenizer_options["eod_token_id"] = eod_token_id
60 if pad_token_id is not None:
61 tokenizer_options["pad_token_id"] = pad_token_id
63 self.unique_identifiers = OrderedDict()
64 self.unique_identifiers["class"] = type(self).__name__
65 self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths)
66 for option, value in tokenizer_options.items():
67 self.unique_identifiers[option] = str(value)
68 self.unique_description = json.dumps(self.unique_identifiers, indent=4)
69 if vocab_size is not None and eod_token_id is not None:
70 self.vocab_size = vocab_size
71 self.eod = eod_token_id
72 self.eos_token_id = eod_token_id
73 self.pad_token_id = pad_token_id
75 def __len__(self) -> int:
76 """Return the configured tokenizer vocabulary size."""
77 return self.vocab_size
79 @property
80 def eos(self) -> int:
81 """Return the end-of-sequence token ID."""
82 return self.eos_token_id
84 @property
85 def pad(self) -> int:
86 """Return the padding token ID."""
87 if self.pad_token_id is None:
88 raise AttributeError("The tokenizer does not define a padding token ID")
90 return self.pad_token_id
93class _HFAutoTokenizer(_DatasetTokenizer):
94 """Wrap a Hugging Face tokenizer with the Dataset tokenizer contract."""
96 def __init__(
97 self,
98 pretrained_model_name_or_path: str,
99 *args: Any,
100 trust_remote_code: bool,
101 **kwargs: Any,
102 ) -> None:
103 """Load a Hugging Face tokenizer and expose the Dataset identity contract.
105 Args:
106 pretrained_model_name_or_path: Model identifier or local tokenizer path.
107 *args: Extra positional arguments forwarded to ``HFAutoTokenizer.from_pretrained``.
108 trust_remote_code: Whether to trust remote code when loading the tokenizer.
109 **kwargs: Extra keyword arguments forwarded to ``HFAutoTokenizer.from_pretrained``.
111 Raises:
112 ValueError: If the tokenizer defines neither ``eod`` nor ``eos_token_id``.
113 """
114 tokenizer_options = dict(kwargs)
115 tokenizer_options["trust_remote_code"] = trust_remote_code
116 if args:
117 tokenizer_options["inputs"] = args
118 super().__init__(pretrained_model_name_or_path, **tokenizer_options)
120 # Transformers is optional and only required when constructing an HF tokenizer.
121 from transformers import AutoTokenizer as HFAutoTokenizer # pylint: disable=C0415
123 self.tokenizer = HFAutoTokenizer.from_pretrained(
124 pretrained_model_name_or_path, *args, trust_remote_code=trust_remote_code, **kwargs
125 )
126 if not hasattr(self.tokenizer, "eod"):
127 eos_token_id = getattr(self.tokenizer, "eos_token_id", None)
128 if eos_token_id is None:
129 raise ValueError("The tokenizer must define either 'eod' or 'eos_token_id'")
130 self.tokenizer.eod = eos_token_id
132 def __getattr__(self, name: str) -> Any:
133 """Delegate tokenizer attributes to the wrapped Hugging Face tokenizer."""
134 return getattr(self.tokenizer, name)
136 def __setattr__(self, name: str, value: Any) -> None:
137 """Delegate runtime tokenizer options after the wrapped tokenizer exists."""
138 tokenizer = self.__dict__.get("tokenizer")
139 if tokenizer is None or name in {"tokenizer", "unique_identifiers", "unique_description"}:
140 object.__setattr__(self, name, value)
141 return
142 setattr(tokenizer, name, value)
144 def __len__(self) -> int:
145 """Return the wrapped tokenizer vocabulary size."""
146 return len(self.tokenizer)
148 def __call__(self, *args: Any, **kwargs: Any) -> Any:
149 """Forward tokenization calls to the wrapped tokenizer."""
150 return self.tokenizer(*args, **kwargs)
153class AutoTokenizer:
154 """
155 Auto tokenizer class that dispatches to appropriate tokenizer implementations.
157 Similar to HuggingFace's AutoTokenizer, but with a custom registry for specialized
158 tokenizer implementations.
160 The dispatch logic is:
161 1. If a custom tokenizer is registered for the model type, use it
162 2. Otherwise, fall back to NeMoAutoTokenizerWithBosEosEnforced
164 Example:
165 >>> # Will use MistralCommonBackend if available for Mistral models
166 >>> tokenizer = NeMoAutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
168 >>> # Force using HF AutoTokenizer with BOS/EOS enforcement
169 >>> tokenizer = NeMoAutoTokenizer.from_pretrained("gpt2", force_default=True)
170 """
172 # Make registry accessible at class level
173 _registry = None
175 @classmethod
176 def from_pretrained(
177 cls,
178 pretrained_model_name_or_path: str | None,
179 *args: Any,
180 force_default: bool = False,
181 force_hf: bool = True,
182 trust_remote_code: bool = False,
183 tokenizer_type: str | None = None,
184 vocab_size: int | None = None,
185 eod_token_id: int | None = None,
186 pad_token_id: int | None = None,
187 **kwargs: Any,
188 ) -> Any:
189 """
190 Load a tokenizer from a pretrained model.
192 Args:
193 pretrained_model_name_or_path: Tokenizer path used for loading or stable Dataset cache identity.
194 *args: Additional positional arguments passed to the tokenizer's from_pretrained.
195 force_default: If True, always use NeMoAutoTokenizerWithBosEosEnforced
196 force_hf: If True, build a wrapped Hugging Face tokenizer.
197 trust_remote_code: Whether to trust remote code when loading config
198 tokenizer_type: Use ``hf`` to load a Hugging Face tokenizer or ``pretokenized`` for metadata only.
199 vocab_size: Pretokenized-corpus vocabulary size.
200 eod_token_id: Pretokenized-corpus end-of-document token ID.
201 pad_token_id: Pretokenized-corpus padding token ID, when one exists.
202 **kwargs: Additional arguments passed to the tokenizer's from_pretrained
204 Returns:
205 A tokenizer instance appropriate for the model type
206 """
207 # ``force_default`` is retained for NeMo API compatibility; only the HF backend is supported.
208 del force_default
209 if tokenizer_type == "pretokenized":
210 if pretrained_model_name_or_path is None:
211 raise ValueError("pretrained_model_name_or_path is required for a pretokenized Dataset")
212 tokenizer = _DatasetTokenizer(
213 pretrained_model_name_or_path,
214 vocab_size=vocab_size,
215 eod_token_id=eod_token_id,
216 pad_token_id=pad_token_id,
217 )
218 elif tokenizer_type not in (None, "hf"):
219 raise ValueError(f"Unsupported tokenizer_type: {tokenizer_type!r}")
220 else:
221 if vocab_size is not None or eod_token_id is not None or pad_token_id is not None:
222 raise ValueError("vocab_size, eod_token_id, and pad_token_id require tokenizer_type='pretokenized'")
223 if pretrained_model_name_or_path is None:
224 raise ValueError("pretrained_model_name_or_path is required for a Hugging Face tokenizer")
225 if tokenizer_type != "hf" and not force_hf:
226 raise ValueError("Only the Hugging Face tokenizer backend is currently supported")
228 # Explicit ``hf`` and the legacy ``force_hf`` switch select the same backend.
229 tokenizer = _HFAutoTokenizer(
230 pretrained_model_name_or_path,
231 *args,
232 trust_remote_code=trust_remote_code,
233 **kwargs,
234 )
236 backend = getattr(tokenizer, "tokenizer", tokenizer)
237 logger.debug(
238 "Built Dataset tokenizer: type=%s, backend=%s, path=%s, vocab_size=%d, "
239 "eod_token_id=%s, eos_token_id=%s, pad_token_id=%s",
240 type(tokenizer).__name__, type(backend).__name__, pretrained_model_name_or_path, len(tokenizer),
241 getattr(tokenizer, "eod", None), getattr(tokenizer, "eos_token_id", None),
242 getattr(tokenizer, "pad_token_id", None),
243 )
244 return tokenizer
247def build_tokenizer(pretrained_model_name_or_path: str, **kwargs: Any) -> Any:
248 """Build the configured LLM tokenizer.
250 Args:
251 pretrained_model_name_or_path: Model identifier or local tokenizer path.
252 **kwargs: Options forwarded to ``AutoTokenizer.from_pretrained``.
254 Returns:
255 The constructed tokenizer.
256 """
257 return AutoTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)