diff --git a/tests/test_encoding.py b/tests/test_encoding.py index b77ca135..4e6bfa1c 100644 --- a/tests/test_encoding.py +++ b/tests/test_encoding.py @@ -110,6 +110,21 @@ def test_encode_surrogate_pairs(): assert enc.encode("\ud83d") == enc.encode("�") +def test_encode_with_unstable_surrogate_pairs(): + enc = tiktoken.get_encoding("cl100k_base") + + # encode_with_unstable must not raise UnicodeEncodeError for surrogate pairs + stable, completions = enc.encode_with_unstable("👍") + stable_via_surrogate, completions_via_surrogate = enc.encode_with_unstable("\ud83d\udc4d") + assert stable == stable_via_surrogate + assert completions == completions_via_surrogate + + # lone surrogate treated the same as the replacement character + stable_lone, _ = enc.encode_with_unstable("\ud83d") + stable_replacement, _ = enc.encode_with_unstable("�") + assert stable_lone == stable_replacement + + @pytest.mark.parametrize("make_enc", ENCODING_FACTORIES) def test_catastrophically_repetitive(make_enc: Callable[[], tiktoken.Encoding]): enc = make_enc() diff --git a/tiktoken/core.py b/tiktoken/core.py index 530f8f59..482e4bf7 100644 --- a/tiktoken/core.py +++ b/tiktoken/core.py @@ -240,7 +240,12 @@ def encode_with_unstable( if match := _special_token_regex(disallowed_special).search(text): raise_disallowed_special_token(match.group()) - return self._core_bpe.encode_with_unstable(text, allowed_special) + try: + return self._core_bpe.encode_with_unstable(text, allowed_special) + except UnicodeEncodeError: + # see comment in encode; same surrogate-pair fixup applied here + text = text.encode("utf-16", "surrogatepass").decode("utf-16", "replace") + return self._core_bpe.encode_with_unstable(text, allowed_special) def encode_single_token(self, text_or_bytes: str | bytes) -> int: """Encodes text corresponding to a single token to its token value.