Repository: explosion/curated-tokenizers (language: python)

Issue:
Make all the piece processors picklable
<!--- Provide a general summary of your changes in the title. -->

## Description

<!--- Use this section to describe your changes. If your changes required
testing, include information about the testing environment and the tests you
ran. If your test fixes a bug reported in an issue, don't forget to include the
issue number. If your PR is still a work in progress, that's totally fine – just
include a note to let us know. -->

The piece processors are not pure Python classes, so could not be pickled. Implement `__reduce__` on all the processors to make them pickable.

This is one step in making spaCy multiprocessing work.

See: https://github.com/explosion/spaCy/issues/13233

### Types of change

<!-- What type of change does your PR cover? Is it a bug fix, an enhancement
or new feature, or a change to the documentation? -->

Feature

## Checklist

<!--- Before you submit the PR, go over this checklist and make sure you can
tick off all the boxes. [] -> [x] -->

- [x] I confirm that I have the right to submit this contribution under the project's MIT license.
- [x] I ran the tests, and all new and existing tests passed.
- [x] My changes don't require a change to the documentation, or if they do, I've added all required information.



Discussion / hints:
Note to self: needs to be backported to spaCy 3.x branch.

Fix patch:
diff --git a/curated_tokenizers/_bbpe.pyx b/curated_tokenizers/_bbpe.pyx
index 98a9c70..ca0aad6 100644
--- a/curated_tokenizers/_bbpe.pyx
+++ b/curated_tokenizers/_bbpe.pyx
@@ -165,3 +165,10 @@ cdef class ByteBPEProcessor:
             piece_ids.append(piece_id)
 
         return piece_ids
+
+    def __reduce__(self):
+        return (unpickle_byte_bpe_processor, (self.vocab, self.merges,))
+
+def unpickle_byte_bpe_processor(vocab: Dict[str, int],
+                                merges: List[Tuple[str, str]]) -> ByteBPEProcessor:
+    return ByteBPEProcessor(vocab, merges)
diff --git a/curated_tokenizers/_spp.pyx b/curated_tokenizers/_spp.pyx
index ad1cba9..44d4760 100644
--- a/curated_tokenizers/_spp.pyx
+++ b/curated_tokenizers/_spp.pyx
@@ -201,6 +201,9 @@ cdef class SentencePieceProcessor:
         _check_status(deref(self.spp).Encode(sentence.encode("utf-8"), &text))
         return text
 
+    def __reduce__(self):
+        return (unpickle_sentence_piece_processor, (self.to_protobuf(),))
+
 cdef _check_status(Status status):
     cdef code = <int> status.code()
     if code == <int> kOk:
@@ -211,3 +214,6 @@ cdef _check_status(Status status):
         raise RuntimeError(status.error_message().decode("utf-8"))
     else:
         raise ValueError(status.error_message().decode("utf-8"))
+
+def unpickle_sentence_piece_processor(protobuf: bytes) -> SentencePieceProcessor:
+    return SentencePieceProcessor.from_protobuf(protobuf)
diff --git a/curated_tokenizers/_wordpiece.pyx b/curated_tokenizers/_wordpiece.pyx
index 204ff1c..b80a2f6 100644
--- a/curated_tokenizers/_wordpiece.pyx
+++ b/curated_tokenizers/_wordpiece.pyx
@@ -127,4 +127,10 @@ cdef class WordPieceProcessor:
                 match.piece_id = idx
                 return match
 
-        return match
\ No newline at end of file
+        return match
+
+    def __reduce__(self):
+        return (unpickle_wordpiece_processor, (self.to_list(),))
+
+def unpickle_wordpiece_processor(pieces: List[str]) -> WordPieceProcessor:
+    return WordPieceProcessor(pieces)


Test patch:
diff --git a/curated_tokenizers/tests/test_bbpe_processor.py b/curated_tokenizers/tests/test_bbpe_processor.py
index 46d6de1..b65239b 100644
--- a/curated_tokenizers/tests/test_bbpe_processor.py
+++ b/curated_tokenizers/tests/test_bbpe_processor.py
@@ -1,8 +1,10 @@
-from curated_tokenizers import ByteBPEProcessor
-import json
 from pathlib import Path
+from pickle import dumps, loads
+
 import pytest
 
+from curated_tokenizers import ByteBPEProcessor
+
 
 @pytest.fixture(scope="module")
 def test_dir(request):
@@ -59,6 +61,7 @@ def test_empty_processor():
         "n",
     ]
 
+
 def test_can_decode(toy_processor):
     assert toy_processor.decode_from_ids(EXAMPLE_PIECE_IDS) == EXAMPLE_TEXT
 
@@ -81,3 +84,11 @@ def test_rejects_incorrect_merges(test_dir):
             vocab=test_dir / "robbert-vocab-1000.json",
             merges=test_dir / "incorrect-merges.txt",
         )
+
+
+def test_pickle(toy_processor):
+    serialized = dumps(toy_processor)
+    deserialized = loads(serialized)
+    assert isinstance(deserialized, ByteBPEProcessor)
+    assert deserialized.vocab == toy_processor.vocab
+    assert deserialized.merges == toy_processor.merges
diff --git a/curated_tokenizers/tests/test_sp.py b/curated_tokenizers/tests/test_sp.py
index f1af6f5..31ce835 100644
--- a/curated_tokenizers/tests/test_sp.py
+++ b/curated_tokenizers/tests/test_sp.py
@@ -1,4 +1,6 @@
 from pathlib import Path
+from pickle import dumps, loads
+
 import pytest
 
 from curated_tokenizers import SentencePieceProcessor
@@ -161,3 +163,10 @@ def test_id_to_piece_and_piece_to_id(toy_model):
         toy_model.id_to_piece(-1)
     with pytest.raises(ValueError):
         toy_model.id_to_piece(len(toy_model))
+
+
+def test_pickle(toy_model):
+    serialized = dumps(toy_model)
+    deserialized = loads(serialized)
+    assert isinstance(deserialized, SentencePieceProcessor)
+    assert deserialized.to_protobuf() == toy_model.to_protobuf()
diff --git a/curated_tokenizers/tests/test_word_piece_processor.py b/curated_tokenizers/tests/test_word_piece_processor.py
index 5b6f273..82b8eb1 100644
--- a/curated_tokenizers/tests/test_word_piece_processor.py
+++ b/curated_tokenizers/tests/test_word_piece_processor.py
@@ -1,7 +1,9 @@
-from curated_tokenizers import WordPieceProcessor
 from pathlib import Path
+from pickle import dumps, loads
+
 import pytest
 
+from curated_tokenizers import WordPieceProcessor
 
 EXAMPLE_TOKENS = [
     "voor",
@@ -76,3 +78,10 @@ def test_piece_id_valid(toy_processor):
     assert toy_processor.is_valid_piece_id(0) == True
     assert toy_processor.is_valid_piece_id(len(TOKEN_PIECES)) == False
     assert toy_processor.is_valid_piece_id(len(TOKEN_PIECES) - 1) == True
+
+
+def test_pickle(toy_processor):
+    serialized = dumps(toy_processor)
+    deserialized = loads(serialized)
+    assert isinstance(deserialized, WordPieceProcessor)
+    assert deserialized.to_list() == toy_processor.to_list()