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.11.0

Branches

Tags

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

Clone

HTTPS

Download ZIP

tiktoken/tests

tests/test_encoding.py

263lines · modeblame

RawDownload
Latest commit unavailable.
c373d9b9Shantanu Jain3 years ago1# Note that there are more actual tests, they're just not currently public :-)
2
3from typing import Callable
4
5import hypothesis
6import hypothesis.strategies as st
7import pytest
8
9import tiktoken
10
11from .test_helpers import ENCODING_FACTORIES, MAX_EXAMPLES
12
13
14def test_simple():
15enc = tiktoken.get_encoding("gpt2")
16assert enc.encode("hello world") == [31373, 995]
17assert enc.decode([31373, 995]) == "hello world"
18assert enc.encode("hello <|endoftext|>", allowed_special="all") == [31373, 220, 50256]
19
20enc = tiktoken.get_encoding("cl100k_base")
21assert enc.encode("hello world") == [15339, 1917]
22assert enc.decode([15339, 1917]) == "hello world"
23assert enc.encode("hello <|endoftext|>", allowed_special="all") == [15339, 220, 100257]
24
25for enc_name in tiktoken.list_encoding_names():
26enc = tiktoken.get_encoding(enc_name)
05e66e8dShantanu1 years ago27for token in range(min(10_000, enc.max_token_value - 1)):
c373d9b9Shantanu Jain3 years ago28assert enc.encode_single_token(enc.decode_single_token_bytes(token)) == token
29
30
31def test_simple_repeated():
32enc = tiktoken.get_encoding("gpt2")
33assert enc.encode("0") == [15]
34assert enc.encode("00") == [405]
35assert enc.encode("000") == [830]
36assert enc.encode("0000") == [2388]
37assert enc.encode("00000") == [20483]
38assert enc.encode("000000") == [10535]
39assert enc.encode("0000000") == [24598]
40assert enc.encode("00000000") == [8269]
41assert enc.encode("000000000") == [10535, 830]
42assert enc.encode("0000000000") == [8269, 405]
43assert enc.encode("00000000000") == [8269, 830]
44assert enc.encode("000000000000") == [8269, 2388]
45assert enc.encode("0000000000000") == [8269, 20483]
46assert enc.encode("00000000000000") == [8269, 10535]
47assert enc.encode("000000000000000") == [8269, 24598]
48assert enc.encode("0000000000000000") == [25645]
49assert enc.encode("00000000000000000") == [8269, 10535, 830]
50
51
5818d566Shantanu11 months ago52def test_large_repeated():
53enc = tiktoken.get_encoding("o200k_base")
54
55with pytest.raises(ValueError):
56enc.encode("x" * 1_000_000)
57
58
c373d9b9Shantanu Jain3 years ago59def test_simple_regex():
60enc = tiktoken.get_encoding("cl100k_base")
61assert enc.encode("rer") == [38149]
62assert enc.encode("'rer") == [2351, 81]
63assert enc.encode("today\n ") == [31213, 198, 220]
64assert enc.encode("today\n \n") == [31213, 27907]
65assert enc.encode("today\n \n") == [31213, 14211]
66
67
68def test_basic_encode():
69enc = tiktoken.get_encoding("r50k_base")
70assert enc.encode("hello world") == [31373, 995]
71
72enc = tiktoken.get_encoding("p50k_base")
73assert enc.encode("hello world") == [31373, 995]
74
75enc = tiktoken.get_encoding("cl100k_base")
76assert enc.encode("hello world") == [15339, 1917]
77assert enc.encode(" \x850") == [220, 126, 227, 15]
78
79
80def test_encode_empty():
81enc = tiktoken.get_encoding("r50k_base")
82assert enc.encode("") == []
83
84
85def test_encode_bytes():
86enc = tiktoken.get_encoding("cl100k_base")
87assert enc._encode_bytes(b" \xec\x8b\xa4\xed") == [62085]
4560a889Shantanu1 years ago88for i in range(10):
89bytestring = b"\x80" * i
90assert enc.decode_bytes(enc._encode_bytes(bytestring)) == bytestring
91
92
93@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
94@hypothesis.given(bytestring=st.binary())
5818d566Shantanu11 months ago95@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
4560a889Shantanu1 years ago96def test_hyp_encode_bytes(make_enc: Callable[[], tiktoken.Encoding], bytestring: bytes):
97enc = make_enc()
98assert enc.decode_bytes(enc._encode_bytes(bytestring)) == bytestring
c373d9b9Shantanu Jain3 years ago99
100
101def test_encode_surrogate_pairs():
102enc = tiktoken.get_encoding("cl100k_base")
103
104assert enc.encode("👍") == [9468, 239, 235]
105# surrogate pair gets converted to codepoint
106assert enc.encode("\ud83d\udc4d") == [9468, 239, 235]
107
108# lone surrogate just gets replaced
109assert enc.encode("\ud83d") == enc.encode("�")
110
111
05e66e8dShantanu1 years ago112@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
113def test_catastrophically_repetitive(make_enc: Callable[[], tiktoken.Encoding]):
114enc = make_enc()
115for c in ["^", "0", "a", "'s", " ", "\n"]:
116big_value = c * 10_000
117assert big_value == enc.decode(enc.encode(big_value))
118
119big_value = " " + big_value
120assert big_value == enc.decode(enc.encode(big_value))
121
122big_value = big_value + "\n"
123assert big_value == enc.decode(enc.encode(big_value))
124
125
c373d9b9Shantanu Jain3 years ago126# ====================
127# Roundtrip
128# ====================
129
130
131@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
132def test_basic_roundtrip(make_enc):
133enc = make_enc()
134for value in (
135"hello",
136"hello ",
137"hello ",
138" hello",
139" hello ",
140" hello ",
141"hello world",
142"请考试我的软件！12345",
143):
144assert value == enc.decode(enc.encode(value))
145assert value == enc.decode(enc.encode_ordinary(value))
146
147
148@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
149@hypothesis.given(text=st.text())
5818d566Shantanu11 months ago150@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
c373d9b9Shantanu Jain3 years ago151def test_hyp_roundtrip(make_enc: Callable[[], tiktoken.Encoding], text):
152enc = make_enc()
153
154assert text == enc.decode(enc.encode(text))
155
156
157@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
158def test_single_token_roundtrip(make_enc: Callable[[], tiktoken.Encoding]):
159enc = make_enc()
160
161for token in range(enc.n_vocab):
162try:
163token_bytes = enc.decode_single_token_bytes(token)
164except KeyError:
165continue
166assert enc.encode_single_token(token_bytes) == token
167
168
169# ====================
170# Special tokens
171# ====================
172
173
174def test_special_token():
175enc = tiktoken.get_encoding("cl100k_base")
176
177eot = enc.encode_single_token("<|endoftext|>")
178assert eot == enc.eot_token
179fip = enc.encode_single_token("<|fim_prefix|>")
180fim = enc.encode_single_token("<|fim_middle|>")
181
182text = "<|endoftext|> hello <|fim_prefix|>"
183assert eot not in enc.encode(text, disallowed_special=())
184with pytest.raises(ValueError):
185enc.encode(text)
186with pytest.raises(ValueError):
187enc.encode(text, disallowed_special="all")
188with pytest.raises(ValueError):
189enc.encode(text, disallowed_special={"<|endoftext|>"})
190with pytest.raises(ValueError):
191enc.encode(text, disallowed_special={"<|fim_prefix|>"})
192
193text = "<|endoftext|> hello <|fim_prefix|> there <|fim_middle|>"
194tokens = enc.encode(text, disallowed_special=())
195assert eot not in tokens
196assert fip not in tokens
197assert fim not in tokens
198
199tokens = enc.encode(text, allowed_special="all", disallowed_special=())
200assert eot in tokens
201assert fip in tokens
202assert fim in tokens
203
204tokens = enc.encode(text, allowed_special="all", disallowed_special="all")
205assert eot in tokens
206assert fip in tokens
207assert fim in tokens
208
209tokens = enc.encode(text, allowed_special={"<|fim_prefix|>"}, disallowed_special=())
210assert eot not in tokens
211assert fip in tokens
212assert fim not in tokens
213
214tokens = enc.encode(text, allowed_special={"<|endoftext|>"}, disallowed_special=())
215assert eot in tokens
216assert fip not in tokens
217assert fim not in tokens
218
219tokens = enc.encode(text, allowed_special={"<|fim_middle|>"}, disallowed_special=())
220assert eot not in tokens
221assert fip not in tokens
222assert fim in tokens
223
224
225@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
226@hypothesis.given(text=st.text())
227@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
228def test_hyp_special_ordinary(make_enc, text: str):
229enc = make_enc()
230assert enc.encode_ordinary(text) == enc.encode(text, disallowed_special=())
231
232
233# ====================
234# Batch encoding
235# ====================
236
237
238@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
239def test_batch_encode(make_enc: Callable[[], tiktoken.Encoding]):
240enc = make_enc()
241text1 = "hello world"
242text2 = "goodbye world"
243
244assert enc.encode_batch([text1]) == [enc.encode(text1)]
245assert enc.encode_batch([text1, text2]) == [enc.encode(text1), enc.encode(text2)]
246
247assert enc.encode_ordinary_batch([text1]) == [enc.encode_ordinary(text1)]
248assert enc.encode_ordinary_batch([text1, text2]) == [
249enc.encode_ordinary(text1),
250enc.encode_ordinary(text2),
251]
252
253
254@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
255@hypothesis.given(batch=st.lists(st.text()))
5818d566Shantanu11 months ago256@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
c373d9b9Shantanu Jain3 years ago257def test_hyp_batch_roundtrip(make_enc: Callable[[], tiktoken.Encoding], batch):
258enc = make_enc()
259
5818d566Shantanu11 months ago260encoded = enc.encode_batch(batch, allowed_special="all")
261assert encoded == [enc.encode(t, allowed_special="all") for t in batch]
c373d9b9Shantanu Jain3 years ago262decoded = enc.decode_batch(encoded)
263assert decoded == batch