Coverage for transformer_lens/utilities/tokenize_utils.py: 94%
115 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""tokenize_utils.
3This module contains utility functions related to tokenization
4"""
6from __future__ import annotations
8import os
9from copy import deepcopy
10from typing import Any, Optional
12import einops
13import numpy as np
14import torch
15from datasets.arrow_dataset import Dataset
16from datasets.iterable_dataset import IterableDataset
17from transformers import AutoTokenizer, PreTrainedTokenizerBase
19from transformer_lens.utilities.hf_utils import keep_single_column
20from transformer_lens.utilities.tensors import get_cumsum_along_dim
23def tokenize_and_concatenate(
24 dataset: Dataset | IterableDataset,
25 tokenizer: PreTrainedTokenizerBase,
26 streaming: bool = False,
27 max_length: int = 1024,
28 column_name: str = "text",
29 add_bos_token: bool = True,
30 num_proc: Optional[int] = 10,
31 set_format: bool = True,
32) -> Dataset | IterableDataset:
33 """Tokenize each document, join with token-level EOS between docs, and reshape into ``(batch, sequence_length)`` rows.
35 Useful for training language models on a large text corpus without per-doc
36 truncation or padding. Absolute-position-embedding models also benefit by
37 avoiding early-token bias (e.g. news articles starting with "CNN").
39 Args:
40 dataset: The dataset to tokenize. Accepts both arrow ``Dataset`` and
41 ``IterableDataset`` (e.g. when loaded with ``streaming=True``).
42 tokenizer (PreTrainedTokenizerBase): The tokenizer. Must have ``bos_token_id`` and ``eos_token_id``.
43 streaming (bool, optional): If True, avoids parallelism. Defaults to False.
44 max_length (int, optional): The length of the context window of the sequence. Defaults to 1024.
45 column_name (str, optional): The name of the text column in the dataset. Defaults to 'text'.
46 add_bos_token (bool, optional): Whether to prepend ``bos_token_id`` to each output row. Defaults to True.
47 num_proc (int, optional): Number of processes for parallel tokenization. ``None``
48 runs in-process -- datasets forks a dill-pickling pool for any int, 1 included.
49 Ignored when ``streaming=True``. Defaults to 10.
50 set_format (bool, optional): If True, calls ``set_format(type="torch")`` on the result. Set False
51 for ``IterableDataset`` (which doesn't support format setting); wrap the output in
52 ``(torch.LongTensor(ex["tokens"]) for ex in tokenized_dataset)`` instead. Defaults to True.
54 Returns:
55 Tokenized dataset of token sequences in a single column ``"tokens"``. Returns the same dataset
56 type as the input (``Dataset`` or ``IterableDataset``).
57 """
58 dataset = keep_single_column(dataset, column_name)
59 has_pad_token = tokenizer.pad_token is not None
60 if not has_pad_token:
61 tokenizer.add_special_tokens({"pad_token": "<PAD>"})
62 seq_len = max_length - 1 if add_bos_token else max_length
64 # Long docs legitimately exceed model_max_length; we slice into rows after.
65 _deprecation_warnings_saved = None
66 if hasattr(tokenizer, "deprecation_warnings"): 66 ↛ 72line 66 didn't jump to line 72 because the condition on line 66 was always true
67 _deprecation_warnings_saved = tokenizer.deprecation_warnings.copy()
68 tokenizer.deprecation_warnings[
69 "sequence-length-is-longer-than-the-specified-maximum"
70 ] = False
72 def tokenize_function(examples: Any) -> dict[str, np.ndarray]:
73 text = examples[column_name]
74 assert tokenizer.eos_token is not None, "Tokenizer must have an EOS token."
75 if not text: 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true
76 return {"tokens": np.array([], dtype=np.int64)}
78 # Per-doc tokenization with explicit token-level EOS — string chunking
79 # could cut tokens mid-doc (#1133); add_special_tokens=False prevents
80 # SentencePiece tokenizers from scattering auto-BOS/EOS per call.
81 encoded = tokenizer(text, add_special_tokens=False)["input_ids"]
82 eos_id = tokenizer.eos_token_id
83 pieces: list[np.ndarray] = []
84 for i, row in enumerate(encoded):
85 pieces.append(np.asarray(row, dtype=np.int64))
86 if i < len(encoded) - 1:
87 pieces.append(np.array([eos_id], dtype=np.int64))
88 if not pieces: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true
89 return {"tokens": np.array([], dtype=np.int64)}
90 tokens = np.concatenate(pieces)
91 num_tokens = len(tokens)
93 if num_tokens < seq_len:
94 num_batches = 1
95 tokens = tokens[:seq_len]
96 if len(tokens) < seq_len: 96 ↛ 106line 96 didn't jump to line 106 because the condition on line 96 was always true
97 # Pad with EOS when no native pad token to avoid OOV IDs.
98 padding_id = tokenizer.eos_token_id if not has_pad_token else tokenizer.pad_token_id
99 tokens = np.concatenate(
100 [tokens, np.full(seq_len - len(tokens), padding_id)], axis=0
101 )
102 else:
103 num_batches = num_tokens // seq_len
104 tokens = tokens[: seq_len * num_batches]
106 tokens = einops.rearrange(
107 tokens, "(batch seq) -> batch seq", batch=num_batches, seq=seq_len
108 )
109 if add_bos_token:
110 prefix = np.full((num_batches, 1), tokenizer.bos_token_id)
111 tokens = np.concatenate([prefix, tokens], axis=1)
112 return {"tokens": tokens}
114 try:
115 # IterableDataset.map() rejects `num_proc` outright (even None), so we
116 # spread the kwarg conditionally rather than always passing it.
117 tokenized_dataset = dataset.map(
118 tokenize_function,
119 batched=True,
120 remove_columns=[column_name],
121 **({"num_proc": num_proc} if not streaming else {}),
122 )
123 finally:
124 if _deprecation_warnings_saved is not None: 124 ↛ 127line 124 didn't jump to line 127 because the condition on line 124 was always true
125 tokenizer.deprecation_warnings.clear()
126 tokenizer.deprecation_warnings.update(_deprecation_warnings_saved)
127 if set_format:
128 tokenized_dataset.set_format(type="torch", columns=["tokens"])
129 return tokenized_dataset
132def get_tokenizer_with_bos(tokenizer: PreTrainedTokenizerBase) -> PreTrainedTokenizerBase:
133 """
134 Returns the tokenizer initialized with add_bos_token=True.
135 Such a tokenizer should be set as the default tokenizer because the tokenization of some
136 tokenizers like LlamaTokenizer are different when bos token is automatically/manually
137 prepended.
139 Note: For tokenizers without a BOS token (e.g., T5), this returns the original tokenizer
140 unchanged since add_bos_token=True would fail in transformers v5+ when bos_token is None.
142 Args:
143 tokenizer (PreTrainedTokenizerBase): The tokenizer to initialize with add_bos_token=True.
145 Returns:
146 PreTrainedTokenizerBase: The tokenizer initialized with add_bos_token=True,
147 or the original tokenizer if it has no BOS token.
148 """
149 # If the tokenizer has no BOS token, we can't set add_bos_token=True
150 # This is the case for T5 and other encoder-decoder models
151 if tokenizer.bos_token is None:
152 return tokenizer
154 init_kwargs = deepcopy(tokenizer.init_kwargs)
155 pretrained_model_name_or_path = init_kwargs.pop("name_or_path")
156 add_bos_token = init_kwargs.pop("add_bos_token", None)
157 if add_bos_token is None:
158 add_bos_token = getattr(tokenizer, "add_bos_token", False)
160 if add_bos_token:
161 tokenizer_with_bos = tokenizer
162 else:
163 huggingface_token = os.environ.get("HF_TOKEN", "")
164 try:
165 tokenizer_with_bos = AutoTokenizer.from_pretrained(
166 pretrained_model_name_or_path,
167 add_bos_token=True,
168 token=huggingface_token if len(huggingface_token) > 0 else None,
169 **init_kwargs,
170 )
171 except ValueError:
172 # tokenizers' TemplateProcessing cannot express special tokens
173 # containing ':' (e.g. Seed-OSS's <seed:bos>), so add_bos_token=True
174 # crashes when rebuilding the post-processor. Keep the original
175 # tokenizer — such models do not auto-prepend BOS anyway.
176 return tokenizer
177 # Preserve padding_side from the original tokenizer, since AutoTokenizer.from_pretrained
178 # resets it to the HuggingFace default (usually "right"). Without this, callers who
179 # explicitly set tokenizer.padding_side = "left" before passing the tokenizer in would
180 # have that setting silently discarded. See issue #801.
181 tokenizer_with_bos.padding_side = tokenizer.padding_side
183 return tokenizer_with_bos
186def get_input_with_manually_prepended_bos(
187 bos_token: Optional[str], input: str | list[str]
188) -> str | list[str]:
189 """
190 Manually prepends the bos token to the input.
192 Args:
193 bos_token (Optional[str]): The BOS token to prepend, or None for a tokenizer
194 that has none (e.g. BERT, T5).
195 input (str | list[str]): The input to prepend the bos token to.
197 Returns:
198 str | list[str]: The input with the bos token manually prepended, or unchanged
199 when there is no BOS token to prepend.
200 """
201 if bos_token is None:
202 # Nothing to prepend. Callers reach this when prepend_bos is asked for and
203 # cfg.tokenizer_prepends_bos is False — correctly so for a BOS-less tokenizer,
204 # since detect_tokenizer_bos_eos() requires a bos_token_id. Concatenating
205 # would raise a TypeError naming neither the tokenizer nor the flag.
206 return input
208 if isinstance(input, str):
209 input = bos_token + input
210 else:
211 input = [bos_token + string for string in input]
212 return input
215def get_tokens_with_bos_removed(
216 tokenizer: PreTrainedTokenizerBase,
217 tokens: torch.Tensor,
218 padding_side: str | None = None,
219) -> torch.Tensor:
220 """
221 Removes the bos token from the beginning of each sequence in `tokens`.
222 The last dimension of `tokens` must be the sequence length.
224 Args:
225 tokenizer (PreTrainedTokenizerBase): The tokenizer used to tokenize the input.
226 tokens (torch.Tensor): The tokenized input.
227 padding_side: The side used to pad ``tokens``. Defaults to the tokenizer setting.
229 Returns:
230 torch.Tensor: The tokenized input with the bos token removed.
231 """
232 if tokenizer.bos_token_id is None:
233 # Nothing to remove (#1628). Callers reach this when cfg.tokenizer_prepends_bos
234 # says the tokenizer prepends a BOS but the tokenizer has none — a stale
235 # flag, since detect_tokenizer_bos_eos() requires a bos_token_id. Trusting
236 # it here would drop a real first token under right padding ([CLS] for a
237 # BERT tokenizer), and compare tokens against None under left padding.
238 return tokens
240 if padding_side is None:
241 padding_side = tokenizer.padding_side
243 if padding_side == "right":
244 return tokens[..., 1:]
246 else:
247 bos_removed_shape = list(tokens.shape)
248 bos_removed_shape[-1] -= 1
250 if tokenizer.bos_token_id == tokenizer.pad_token_id:
251 is_not_pad_token = tokens.ne(tokenizer.pad_token_id)
252 is_leading_pad = get_cumsum_along_dim(is_not_pad_token, -1, reverse=False) == 0
253 real_bos_positions = is_leading_pad.sum(-1) - 1
254 else:
255 real_bos_positions = (tokens == tokenizer.bos_token_id).int().argmax(-1)
257 tokens = tokens.scatter(dim=1, index=real_bos_positions.unsqueeze(-1), value=-100)
258 return tokens[tokens != -100].view(*bos_removed_shape)
261def get_attention_mask(
262 tokenizer: PreTrainedTokenizerBase,
263 tokens: torch.Tensor,
264 prepend_bos: bool,
265 padding_side: str | None = None,
266) -> torch.Tensor:
267 """
268 Computes the attention mask for the tokenized input.
269 NOTE: Only the leftmost leading pads (when `padding_side == left`)
270 or rightmost trailing pads (when `padding_side == right`) are
271 considered as real pad tokens that should not be attended.
273 Args:
274 tokenizer (PreTrainedTokenizerBase): The tokenizer used for tokenization.
275 tokens (torch.Tensor): The tokenized input.
276 prepend_bos (bool): If True, a BOS token is prepended to the input.
277 padding_side: The side used to pad ``tokens``. Defaults to the tokenizer setting.
279 Returns:
280 torch.Tensor: The attention mask for the input.
281 """
283 # Initialize the attention mask with ones (indicating all tokens should be attended to)
284 attention_mask = torch.ones_like(tokens)
285 if tokenizer is None: 285 ↛ 286line 285 didn't jump to line 286 because the condition on line 285 was never true
286 return attention_mask
287 if padding_side is None:
288 padding_side = tokenizer.padding_side
289 is_not_pad_token = tokens.ne(tokenizer.pad_token_id)
291 if padding_side == "right":
292 # Zero-out the rightmost trailing pad tokens
293 is_trailing_pad = get_cumsum_along_dim(is_not_pad_token, -1, reverse=True) == 0
294 attention_mask[is_trailing_pad] = 0
295 else:
296 # Zero-out the leftmost leading pad tokens
297 is_leading_pad = get_cumsum_along_dim(is_not_pad_token, -1, reverse=False) == 0
298 attention_mask[is_leading_pad] = 0
300 # Unmask BOS when it shares the same ID as pad token
301 if prepend_bos and tokenizer.bos_token_id == tokenizer.pad_token_id:
302 pad_bos_positions = is_leading_pad.sum(-1) - 1
303 attention_mask[torch.arange(attention_mask.shape[0]), pad_bos_positions] = 1
305 return attention_mask