GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/tiktoken

Public

mirrored from https://github.com/openai/tiktokenAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
0.6.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

tiktoken/tiktoken

tiktoken/core.py

407lines · modecode

RawDownload
Latest commit unavailable.
1from __future__ import annotations
2
3import functools
4from concurrent.futures import ThreadPoolExecutor
5from typing import AbstractSet, Collection, Literal, NoReturn, Optional, Union
6
7import regex
8
9from tiktoken import _tiktoken
10
11
12class Encoding:
13 def __init__(
14 self,
15 name: str,
16 *,
17 pat_str: str,
18 mergeable_ranks: dict[bytes, int],
19 special_tokens: dict[str, int],
20 explicit_n_vocab: Optional[int] = None,
21 ):
22 """Creates an Encoding object.
23
24 See openai_public.py for examples of how to construct an Encoding object.
25
26 Args:
27 name: The name of the encoding. It should be clear from the name of the encoding
28 what behaviour to expect, in particular, encodings with different special tokens
29 should have different names.
30 pat_str: A regex pattern string that is used to split the input text.
31 mergeable_ranks: A dictionary mapping mergeable token bytes to their ranks. The ranks
32 must correspond to merge priority.
33 special_tokens: A dictionary mapping special token strings to their token values.
34 explicit_n_vocab: The number of tokens in the vocabulary. If provided, it is checked
35 that the number of mergeable tokens and special tokens is equal to this number.
36 """
37 self.name = name
38
39 self._pat_str = pat_str
40 self._mergeable_ranks = mergeable_ranks
41 self._special_tokens = special_tokens
42
43 self.max_token_value = max(
44 max(mergeable_ranks.values()), max(special_tokens.values(), default=0)
45 )
46 if explicit_n_vocab:
47 assert len(mergeable_ranks) + len(special_tokens) == explicit_n_vocab
48 assert self.max_token_value == explicit_n_vocab - 1
49
50 self._core_bpe = _tiktoken.CoreBPE(mergeable_ranks, special_tokens, pat_str)
51
52 def __repr__(self) -> str:
53 return f"<Encoding {self.name!r}>"
54
55 # ====================
56 # Encoding
57 # ====================
58
59 def encode_ordinary(self, text: str) -> list[int]:
60 """Encodes a string into tokens, ignoring special tokens.
61
62 This is equivalent to `encode(text, disallowed_special=())` (but slightly faster).
63
64 ```
65 >>> enc.encode_ordinary("hello world")
66 [31373, 995]
67 """
68 try:
69 return self._core_bpe.encode_ordinary(text)
70 except UnicodeEncodeError:
71 # See comment in encode
72 text = text.encode("utf-16", "surrogatepass").decode("utf-16", "replace")
73 return self._core_bpe.encode_ordinary(text)
74
75 def encode(
76 self,
77 text: str,
78 *,
79 allowed_special: Union[Literal["all"], AbstractSet[str]] = set(), # noqa: B006
80 disallowed_special: Union[Literal["all"], Collection[str]] = "all",
81 ) -> list[int]:
82 """Encodes a string into tokens.
83
84 Special tokens are artificial tokens used to unlock capabilities from a model,
85 such as fill-in-the-middle. So we want to be careful about accidentally encoding special
86 tokens, since they can be used to trick a model into doing something we don't want it to do.
87
88 Hence, by default, encode will raise an error if it encounters text that corresponds
89 to a special token. This can be controlled on a per-token level using the `allowed_special`
90 and `disallowed_special` parameters. In particular:
91 - Setting `disallowed_special` to () will prevent this function from raising errors and
92 cause all text corresponding to special tokens to be encoded as natural text.
93 - Setting `allowed_special` to "all" will cause this function to treat all text
94 corresponding to special tokens to be encoded as special tokens.
95
96 ```
97 >>> enc.encode("hello world")
98 [31373, 995]
99 >>> enc.encode("<|endoftext|>", allowed_special={"<|endoftext|>"})
100 [50256]
101 >>> enc.encode("<|endoftext|>", allowed_special="all")
102 [50256]
103 >>> enc.encode("<|endoftext|>")
104 # Raises ValueError
105 >>> enc.encode("<|endoftext|>", disallowed_special=())
106 [27, 91, 437, 1659, 5239, 91, 29]
107 ```
108 """
109 if allowed_special == "all":
110 allowed_special = self.special_tokens_set
111 if disallowed_special == "all":
112 disallowed_special = self.special_tokens_set - allowed_special
113 if disallowed_special:
114 if not isinstance(disallowed_special, frozenset):
115 disallowed_special = frozenset(disallowed_special)
116 if match := _special_token_regex(disallowed_special).search(text):
117 raise_disallowed_special_token(match.group())
118
119 # https://github.com/PyO3/pyo3/pull/3632
120 if isinstance(allowed_special, frozenset):
121 allowed_special = set(allowed_special)
122
123 try:
124 return self._core_bpe.encode(text, allowed_special)
125 except UnicodeEncodeError:
126 # BPE operates on bytes, but the regex operates on unicode. If we pass a str that is
127 # invalid UTF-8 to Rust, it will rightfully complain. Here we do a quick and dirty
128 # fixup for any surrogate pairs that may have sneaked their way into the text.
129 # Technically, this introduces a place where encode + decode doesn't roundtrip a Python
130 # string, but given that this is input we want to support, maybe that's okay.
131 # Also we use errors="replace" to handle weird things like lone surrogates.
132 text = text.encode("utf-16", "surrogatepass").decode("utf-16", "replace")
133 return self._core_bpe.encode(text, allowed_special)
134
135 def encode_ordinary_batch(self, text: list[str], *, num_threads: int = 8) -> list[list[int]]:
136 """Encodes a list of strings into tokens, in parallel, ignoring special tokens.
137
138 This is equivalent to `encode_batch(text, disallowed_special=())` (but slightly faster).
139
140 ```
141 >>> enc.encode_ordinary_batch(["hello world", "goodbye world"])
142 [[31373, 995], [11274, 16390, 995]]
143 ```
144 """
145 encoder = functools.partial(self.encode_ordinary)
146 with ThreadPoolExecutor(num_threads) as e:
147 return list(e.map(encoder, text))
148
149 def encode_batch(
150 self,
151 text: list[str],
152 *,
153 num_threads: int = 8,
154 allowed_special: Union[Literal["all"], AbstractSet[str]] = set(), # noqa: B006
155 disallowed_special: Union[Literal["all"], Collection[str]] = "all",
156 ) -> list[list[int]]:
157 """Encodes a list of strings into tokens, in parallel.
158
159 See `encode` for more details on `allowed_special` and `disallowed_special`.
160
161 ```
162 >>> enc.encode_batch(["hello world", "goodbye world"])
163 [[31373, 995], [11274, 16390, 995]]
164 ```
165 """
166 if allowed_special == "all":
167 allowed_special = self.special_tokens_set
168 if disallowed_special == "all":
169 disallowed_special = self.special_tokens_set - allowed_special
170 if not isinstance(disallowed_special, frozenset):
171 disallowed_special = frozenset(disallowed_special)
172
173 encoder = functools.partial(
174 self.encode, allowed_special=allowed_special, disallowed_special=disallowed_special
175 )
176 with ThreadPoolExecutor(num_threads) as e:
177 return list(e.map(encoder, text))
178
179 def encode_with_unstable(
180 self,
181 text: str,
182 *,
183 allowed_special: Union[Literal["all"], AbstractSet[str]] = set(), # noqa: B006
184 disallowed_special: Union[Literal["all"], Collection[str]] = "all",
185 ) -> tuple[list[int], list[list[int]]]:
186 """Encodes a string into stable tokens and possible completion sequences.
187
188 Note that the stable tokens will only represent a substring of `text`.
189
190 See `encode` for more details on `allowed_special` and `disallowed_special`.
191
192 This API should itself be considered unstable.
193
194 ```
195 >>> enc.encode_with_unstable("hello fanta")
196 ([31373], [(277, 4910), (5113, 265), ..., (8842,)])
197
198 >>> text = "..."
199 >>> stable_tokens, completions = enc.encode_with_unstable(text)
200 >>> assert text.encode().startswith(enc.decode_bytes(stable_tokens))
201 >>> assert all(enc.decode_bytes(stable_tokens + seq).startswith(text.encode()) for seq in completions)
202 ```
203 """
204 if allowed_special == "all":
205 allowed_special = self.special_tokens_set
206 if disallowed_special == "all":
207 disallowed_special = self.special_tokens_set - allowed_special
208 if disallowed_special:
209 if not isinstance(disallowed_special, frozenset):
210 disallowed_special = frozenset(disallowed_special)
211 if match := _special_token_regex(disallowed_special).search(text):
212 raise_disallowed_special_token(match.group())
213
214 return self._core_bpe.encode_with_unstable(text, allowed_special)
215
216 def encode_single_token(self, text_or_bytes: Union[str, bytes]) -> int:
217 """Encodes text corresponding to a single token to its token value.
218
219 NOTE: this will encode all special tokens.
220
221 Raises `KeyError` if the token is not in the vocabulary.
222
223 ```
224 >>> enc.encode_single_token("hello")
225 31373
226 ```
227 """
228 if isinstance(text_or_bytes, str):
229 text_or_bytes = text_or_bytes.encode("utf-8")
230 return self._core_bpe.encode_single_token(text_or_bytes)
231
232 # ====================
233 # Decoding
234 # ====================
235
236 def decode_bytes(self, tokens: list[int]) -> bytes:
237 """Decodes a list of tokens into bytes.
238
239 ```
240 >>> enc.decode_bytes([31373, 995])
241 b'hello world'
242 ```
243 """
244 return self._core_bpe.decode_bytes(tokens)
245
246 def decode(self, tokens: list[int], errors: str = "replace") -> str:
247 """Decodes a list of tokens into a string.
248
249 WARNING: the default behaviour of this function is lossy, since decoded bytes are not
250 guaranteed to be valid UTF-8. You can control this behaviour using the `errors` parameter,
251 for instance, setting `errors=strict`.
252
253 ```
254 >>> enc.decode([31373, 995])
255 'hello world'
256 ```
257 """
258 return self._core_bpe.decode_bytes(tokens).decode("utf-8", errors=errors)
259
260 def decode_single_token_bytes(self, token: int) -> bytes:
261 """Decodes a token into bytes.
262
263 NOTE: this will decode all special tokens.
264
265 Raises `KeyError` if the token is not in the vocabulary.
266
267 ```
268 >>> enc.decode_single_token_bytes(31373)
269 b'hello'
270 ```
271 """
272 return self._core_bpe.decode_single_token_bytes(token)
273
274 def decode_tokens_bytes(self, tokens: list[int]) -> list[bytes]:
275 """Decodes a list of tokens into a list of bytes.
276
277 Useful for visualising tokenisation.
278 >>> enc.decode_tokens_bytes([31373, 995])
279 [b'hello', b' world']
280 """
281 return [self.decode_single_token_bytes(token) for token in tokens]
282
283 def decode_with_offsets(self, tokens: list[int]) -> tuple[str, list[int]]:
284 """Decodes a list of tokens into a string and a list of offsets.
285
286 Each offset is the index into text corresponding to the start of each token.
287 If UTF-8 character boundaries do not line up with token boundaries, the offset is the index
288 of the first character that contains bytes from the token.
289
290 This will currently raise if given tokens that decode to invalid UTF-8; this behaviour may
291 change in the future to be more permissive.
292
293 >>> enc.decode_with_offsets([31373, 995])
294 ('hello world', [0, 5])
295 """
296 token_bytes = self.decode_tokens_bytes(tokens)
297
298 text_len = 0
299 offsets = []
300 for token in token_bytes:
301 offsets.append(max(0, text_len - (0x80 <= token[0] < 0xC0)))
302 text_len += sum(1 for c in token if not 0x80 <= c < 0xC0)
303
304 # TODO: assess correctness for errors="ignore" and errors="replace"
305 text = b"".join(token_bytes).decode("utf-8", errors="strict")
306 return text, offsets
307
308 def decode_batch(
309 self, batch: list[list[int]], *, errors: str = "replace", num_threads: int = 8
310 ) -> list[str]:
311 """Decodes a batch (list of lists of tokens) into a list of strings."""
312 decoder = functools.partial(self.decode, errors=errors)
313 with ThreadPoolExecutor(num_threads) as e:
314 return list(e.map(decoder, batch))
315
316 def decode_bytes_batch(self, batch: list[list[int]], *, num_threads: int = 8) -> list[bytes]:
317 """Decodes a batch (list of lists of tokens) into a list of bytes."""
318 with ThreadPoolExecutor(num_threads) as e:
319 return list(e.map(self.decode_bytes, batch))
320
321 # ====================
322 # Miscellaneous
323 # ====================
324
325 def token_byte_values(self) -> list[bytes]:
326 """Returns the list of all token byte values."""
327 return self._core_bpe.token_byte_values()
328
329 @property
330 def eot_token(self) -> int:
331 return self._special_tokens["<|endoftext|>"]
332
333 @functools.cached_property
334 def special_tokens_set(self) -> set[str]:
335 return set(self._special_tokens.keys())
336
337 @property
338 def n_vocab(self) -> int:
339 """For backwards compatibility. Prefer to use `enc.max_token_value + 1`."""
340 return self.max_token_value + 1
341
342 # ====================
343 # Private
344 # ====================
345
346 def _encode_single_piece(self, text_or_bytes: Union[str, bytes]) -> list[int]:
347 """Encodes text corresponding to bytes without a regex split.
348
349 NOTE: this will not encode any special tokens.
350
351 ```
352 >>> enc.encode_single_piece("helloqqqq")
353 [31373, 38227, 38227]
354 ```
355 """
356 if isinstance(text_or_bytes, str):
357 text_or_bytes = text_or_bytes.encode("utf-8")
358 return self._core_bpe.encode_single_piece(text_or_bytes)
359
360 def _encode_only_native_bpe(self, text: str) -> list[int]:
361 """Encodes a string into tokens, but do regex splitting in Python."""
362 _unused_pat = regex.compile(self._pat_str)
363 ret = []
364 for piece in regex.findall(_unused_pat, text):
365 ret.extend(self._core_bpe.encode_single_piece(piece))
366 return ret
367
368 def _encode_bytes(self, text: bytes) -> list[int]:
369 return self._core_bpe._encode_bytes(text)
370
371 def __getstate__(self) -> object:
372 import tiktoken.registry
373
374 # As an optimisation, pickle registered encodings by reference
375 if self is tiktoken.registry.ENCODINGS.get(self.name):
376 return self.name
377 return {
378 "name": self.name,
379 "pat_str": self._pat_str,
380 "mergeable_ranks": self._mergeable_ranks,
381 "special_tokens": self._special_tokens,
382 }
383
384 def __setstate__(self, value: object) -> None:
385 import tiktoken.registry
386
387 if isinstance(value, str):
388 self.__dict__ = tiktoken.registry.get_encoding(value).__dict__
389 return
390 self.__init__(**value)
391
392
393@functools.lru_cache(maxsize=128)
394def _special_token_regex(tokens: frozenset[str]) -> "regex.Pattern[str]":
395 inner = "|".join(regex.escape(token) for token in tokens)
396 return regex.compile(f"({inner})")
397
398
399def raise_disallowed_special_token(token: str) -> NoReturn:
400 raise ValueError(
401 f"Encountered text corresponding to disallowed special token {token!r}.\n"
402 "If you want this text to be encoded as a special token, "
403 f"pass it to `allowed_special`, e.g. `allowed_special={{{token!r}, ...}}`.\n"
404 f"If you want this text to be encoded as normal text, disable the check for this token "
405 f"by passing `disallowed_special=(enc.special_tokens_set - {{{token!r}}})`.\n"
406 "To disable this check for all special tokens, pass `disallowed_special=()`.\n"
407 )
408