SentencePieceProcessor: Handle zero-length protobufs during deseria…

…lization
explosion · Jul 4, 2023 · 74288ba · 74288ba
1 parent 4bdaf4c
commit 74288ba
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/curated_tokenizers/_spp.pyx b/curated_tokenizers/_spp.pyx
@@ -42,6 +42,8 @@ cdef class SentencePieceProcessor:
                 protocol buffer.
         """
         cdef SentencePieceProcessor processor = SentencePieceProcessor.__new__(SentencePieceProcessor)
+        if len(protobuf) == 0:
+            return processor
         cdef string_view protobuf_view = string_view(protobuf, len(protobuf))
         _check_status(deref(processor.spp).LoadFromSerializedProto(protobuf_view))
         return processor

diff --git a/curated_tokenizers/tests/test_sp.py b/curated_tokenizers/tests/test_sp.py
@@ -1,6 +1,6 @@
 from pathlib import Path
-import pytest
 
+import pytest
 from curated_tokenizers import SentencePieceProcessor
 
 
@@ -22,6 +22,10 @@ def test_load_proto(test_dir):
     serialized_data = spp.to_protobuf()
     assert serialized_data == data
 
+    # Zero-length buffer.
+    spp = SentencePieceProcessor.from_protobuf(bytes())
+    assert spp.to_protobuf() == bytes()
+
 
 def test_load_unknown_file():
     with pytest.raises(OSError, match=r"No such file"):