GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/tiktoken

Public

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

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
subquad2

Branches

Tags

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

Clone

HTTPS

Download ZIP

tiktoken/tests

tests/test_encoding.py

256lines · 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
52def test_simple_regex():
53enc = tiktoken.get_encoding("cl100k_base")
54assert enc.encode("rer") == [38149]
55assert enc.encode("'rer") == [2351, 81]
56assert enc.encode("today\n ") == [31213, 198, 220]
57assert enc.encode("today\n \n") == [31213, 27907]
58assert enc.encode("today\n \n") == [31213, 14211]
59
60
61def test_basic_encode():
62enc = tiktoken.get_encoding("r50k_base")
63assert enc.encode("hello world") == [31373, 995]
64
65enc = tiktoken.get_encoding("p50k_base")
66assert enc.encode("hello world") == [31373, 995]
67
68enc = tiktoken.get_encoding("cl100k_base")
69assert enc.encode("hello world") == [15339, 1917]
70assert enc.encode(" \x850") == [220, 126, 227, 15]
71
72
73def test_encode_empty():
74enc = tiktoken.get_encoding("r50k_base")
75assert enc.encode("") == []
76
77
78def test_encode_bytes():
79enc = tiktoken.get_encoding("cl100k_base")
80assert enc._encode_bytes(b" \xec\x8b\xa4\xed") == [62085]
4560a889Shantanu1 years ago81for i in range(10):
82bytestring = b"\x80" * i
83assert enc.decode_bytes(enc._encode_bytes(bytestring)) == bytestring
84
85
86@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
87@hypothesis.given(bytestring=st.binary())
88@hypothesis.settings(deadline=None)
89def test_hyp_encode_bytes(make_enc: Callable[[], tiktoken.Encoding], bytestring: bytes):
90enc = make_enc()
91assert enc.decode_bytes(enc._encode_bytes(bytestring)) == bytestring
c373d9b9Shantanu Jain3 years ago92
93
94def test_encode_surrogate_pairs():
95enc = tiktoken.get_encoding("cl100k_base")
96
97assert enc.encode("👍") == [9468, 239, 235]
98# surrogate pair gets converted to codepoint
99assert enc.encode("\ud83d\udc4d") == [9468, 239, 235]
100
101# lone surrogate just gets replaced
102assert enc.encode("\ud83d") == enc.encode("�")
103
104
05e66e8dShantanu1 years ago105@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
106def test_catastrophically_repetitive(make_enc: Callable[[], tiktoken.Encoding]):
107enc = make_enc()
108for c in ["^", "0", "a", "'s", " ", "\n"]:
109big_value = c * 10_000
110assert big_value == enc.decode(enc.encode(big_value))
111
112big_value = " " + big_value
113assert big_value == enc.decode(enc.encode(big_value))
114
115big_value = big_value + "\n"
116assert big_value == enc.decode(enc.encode(big_value))
117
118
c373d9b9Shantanu Jain3 years ago119# ====================
120# Roundtrip
121# ====================
122
123
124@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
125def test_basic_roundtrip(make_enc):
126enc = make_enc()
127for value in (
128"hello",
129"hello ",
130"hello ",
131" hello",
132" hello ",
133" hello ",
134"hello world",
135"请考试我的软件！12345",
136):
137assert value == enc.decode(enc.encode(value))
138assert value == enc.decode(enc.encode_ordinary(value))
139
140
141@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
142@hypothesis.given(text=st.text())
143@hypothesis.settings(deadline=None)
144def test_hyp_roundtrip(make_enc: Callable[[], tiktoken.Encoding], text):
145enc = make_enc()
146
147assert text == enc.decode(enc.encode(text))
148
149
150@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
151def test_single_token_roundtrip(make_enc: Callable[[], tiktoken.Encoding]):
152enc = make_enc()
153
154for token in range(enc.n_vocab):
155try:
156token_bytes = enc.decode_single_token_bytes(token)
157except KeyError:
158continue
159assert enc.encode_single_token(token_bytes) == token
160
161
162# ====================
163# Special tokens
164# ====================
165
166
167def test_special_token():
168enc = tiktoken.get_encoding("cl100k_base")
169
170eot = enc.encode_single_token("<|endoftext|>")
171assert eot == enc.eot_token
172fip = enc.encode_single_token("<|fim_prefix|>")
173fim = enc.encode_single_token("<|fim_middle|>")
174
175text = "<|endoftext|> hello <|fim_prefix|>"
176assert eot not in enc.encode(text, disallowed_special=())
177with pytest.raises(ValueError):
178enc.encode(text)
179with pytest.raises(ValueError):
180enc.encode(text, disallowed_special="all")
181with pytest.raises(ValueError):
182enc.encode(text, disallowed_special={"<|endoftext|>"})
183with pytest.raises(ValueError):
184enc.encode(text, disallowed_special={"<|fim_prefix|>"})
185
186text = "<|endoftext|> hello <|fim_prefix|> there <|fim_middle|>"
187tokens = enc.encode(text, disallowed_special=())
188assert eot not in tokens
189assert fip not in tokens
190assert fim not in tokens
191
192tokens = enc.encode(text, allowed_special="all", disallowed_special=())
193assert eot in tokens
194assert fip in tokens
195assert fim in tokens
196
197tokens = enc.encode(text, allowed_special="all", disallowed_special="all")
198assert eot in tokens
199assert fip in tokens
200assert fim in tokens
201
202tokens = enc.encode(text, allowed_special={"<|fim_prefix|>"}, disallowed_special=())
203assert eot not in tokens
204assert fip in tokens
205assert fim not in tokens
206
207tokens = enc.encode(text, allowed_special={"<|endoftext|>"}, disallowed_special=())
208assert eot in tokens
209assert fip not in tokens
210assert fim not in tokens
211
212tokens = enc.encode(text, allowed_special={"<|fim_middle|>"}, disallowed_special=())
213assert eot not in tokens
214assert fip not in tokens
215assert fim in tokens
216
217
218@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
219@hypothesis.given(text=st.text())
220@hypothesis.settings(deadline=None, max_examples=MAX_EXAMPLES)
221def test_hyp_special_ordinary(make_enc, text: str):
222enc = make_enc()
223assert enc.encode_ordinary(text) == enc.encode(text, disallowed_special=())
224
225
226# ====================
227# Batch encoding
228# ====================
229
230
231@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
232def test_batch_encode(make_enc: Callable[[], tiktoken.Encoding]):
233enc = make_enc()
234text1 = "hello world"
235text2 = "goodbye world"
236
237assert enc.encode_batch([text1]) == [enc.encode(text1)]
238assert enc.encode_batch([text1, text2]) == [enc.encode(text1), enc.encode(text2)]
239
240assert enc.encode_ordinary_batch([text1]) == [enc.encode_ordinary(text1)]
241assert enc.encode_ordinary_batch([text1, text2]) == [
242enc.encode_ordinary(text1),
243enc.encode_ordinary(text2),
244]
245
246
247@pytest.mark.parametrize("make_enc", ENCODING_FACTORIES)
248@hypothesis.given(batch=st.lists(st.text()))
249@hypothesis.settings(deadline=None)
250def test_hyp_batch_roundtrip(make_enc: Callable[[], tiktoken.Encoding], batch):
251enc = make_enc()
252
253encoded = enc.encode_batch(batch)
254assert encoded == [enc.encode(t) for t in batch]
255decoded = enc.decode_batch(encoded)
256assert decoded == batch