Handle surrogate pairs in encode_with_unstable

src/lib.rs

@@ -633,12 +633,14 @@ impl CoreBPE {
         let decoder: HashMap<Rank, Vec<u8>> =
             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<Rank, Vec<u8>> = special_tokens_encoder
             .iter()

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={},
+        )

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)]

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.