diff --git a/src/lib.rs b/src/lib.rs index ea54eac8..0315600e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -633,12 +633,14 @@ impl CoreBPE { let decoder: HashMap> = encoder.iter().map(|(k, v)| (*v, k.clone())).collect(); - assert!( - encoder.len() == decoder.len(), - "Encoder and decoder must be of equal length. Encoder length: {}, decoder length: {}.\nMaybe you had duplicate token indices in your encoder?", - encoder.len(), - decoder.len() - ); + if encoder.len() != decoder.len() { + return Err(format!( + "Encoder and decoder must be of equal length. Encoder length: {}, decoder length: {}.\nMaybe you had duplicate token indices in your encoder?", + encoder.len(), + decoder.len() + ) + .into()); + } let special_tokens_decoder: HashMap> = special_tokens_encoder .iter() diff --git a/tests/test_duplicate_ranks.py b/tests/test_duplicate_ranks.py new file mode 100644 index 00000000..306e8171 --- /dev/null +++ b/tests/test_duplicate_ranks.py @@ -0,0 +1,13 @@ +import pytest + +import tiktoken + + +def test_duplicate_mergeable_ranks_raise_value_error(): + with pytest.raises(ValueError): + tiktoken.Encoding( + name="duplicate_ranks", + pat_str=r".", + mergeable_ranks={b"a": 0, b"b": 0, b"c": 1}, + special_tokens={}, + ) diff --git a/tests/test_unstable.py b/tests/test_unstable.py new file mode 100644 index 00000000..7b7f8e28 --- /dev/null +++ b/tests/test_unstable.py @@ -0,0 +1,14 @@ +import tiktoken + + +def test_encode_with_unstable_surrogate_pairs(): + enc = tiktoken.Encoding( + name="test", + pat_str=r"(?s:.)", + mergeable_ranks={bytes([i]): i for i in range(256)}, + special_tokens={}, + ) + + for text in ["py\ud83d\udc4d", "py\ud83d"]: + stable_tokens, _ = enc.encode_with_unstable(text) + assert stable_tokens == enc.encode(text)[: len(stable_tokens)] diff --git a/tiktoken/core.py b/tiktoken/core.py index 530f8f59..953406c2 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 + 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.