diff --git a/tests/test_simple_public.py b/tests/test_simple_public.py index 4410923..ce10da4 100644 --- a/tests/test_simple_public.py +++ b/tests/test_simple_public.py @@ -24,3 +24,5 @@ def test_encoding_for_model(): assert enc.name == "gpt2" enc = tiktoken.encoding_for_model("text-davinci-003") assert enc.name == "p50k_base" + enc = tiktoken.encoding_for_model("text-davinci-edit-001") + assert enc.name == "p50k_edit" diff --git a/tiktoken_ext/openai_public.py b/tiktoken_ext/openai_public.py index a64db9f..522d58f 100644 --- a/tiktoken_ext/openai_public.py +++ b/tiktoken_ext/openai_public.py @@ -84,4 +84,5 @@ ENCODING_CONSTRUCTORS = { "r50k_base": r50k_base, "p50k_base": p50k_base, "cl100k_base": cl100k_base, + "p50k_edit": p50k_edit, }