fix: multiprocessing pickle error with tokenizer

src/mistral_common/tokens/tokenizers/base.py

@@ -2,6 +2,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
 from typing import Generic, List, Optional, Protocol, Tuple, TypeVar, Union
 
 import numpy as np
@@ -236,6 +237,12 @@ def to_string(self, tokens: List[int]) -> str:
     @abstractmethod
     def _to_string(self, tokens: List[int]) -> str: ...
 
+    @property
+    @abstractmethod
+    def file_path(self) -> Path:
+        r"""The file path of the tokenizer."""
+        ...
+
 
 InstructRequestType = TypeVar("InstructRequestType", bound=InstructRequest)
 FIMRequestType = TypeVar("FIMRequestType", bound=FIMRequest)

src/mistral_common/tokens/tokenizers/mistral.py

@@ -1,6 +1,6 @@
 import warnings
 from pathlib import Path
-from typing import Callable, Dict, Generic, List, Optional, Union
+from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, Union
 
 from mistral_common.exceptions import (
     TokenizerException,
@@ -109,6 +109,15 @@ def __init__(
             instruct_tokenizer
         )
 
+    def __reduce__(self) -> Tuple[Callable, Tuple[Any, ...]]:
+        """
+        Provides a recipe for pickling (serializing) this object, which is necessary for use with multiprocessing.
+
+        Returns:
+            A tuple of the factory function and the arguments to reconstruct the object from its source file.
+        """
+        return MistralTokenizer.from_file, (self.instruct_tokenizer.tokenizer.file_path,)
+
     @classmethod
     def _data_path(cls) -> Path:
         return Path(__file__).parents[2] / "data"
@@ -238,7 +247,7 @@ def from_hf_hub(
     @classmethod
     def from_file(
         cls,
-        tokenizer_filename: str,
+        tokenizer_filename: Union[str, Path],
         mode: ValidationMode = ValidationMode.test,
     ) -> "MistralTokenizer":
         r"""Loads a tokenizer from a file.

src/mistral_common/tokens/tokenizers/sentencepiece.py

@@ -28,8 +28,11 @@ def is_sentencepiece(path: Union[str, Path]) -> bool:
     return path.is_file() and any(path.name.endswith(suffix) for suffix in suffixes)
 
 
-def get_spm_version(tokenizer_filename: str, raise_deprecated: bool = False) -> TokenizerVersion:
+def get_spm_version(tokenizer_filename: Union[str, Path], raise_deprecated: bool = False) -> TokenizerVersion:
     r"""Get the version of the tokenizer from the filename."""
+    if isinstance(tokenizer_filename, Path):
+        tokenizer_filename = str(tokenizer_filename)
+
     _version_str = tokenizer_filename.split(".")[-1]
     if _version_str != "model":  # filter tokenizer_filename == "/path/to/tokenizer.model" case
         _version_str = _version_str.split("m")[0]
@@ -47,8 +50,11 @@ def get_spm_version(tokenizer_filename: str, raise_deprecated: bool = False) ->
     return TokenizerVersion(_version_str)
 
 
-def get_mm_config(tokenizer_filename: str) -> Optional[MultimodalConfig]:
+def get_mm_config(tokenizer_filename: Union[str, Path]) -> Optional[MultimodalConfig]:
     r"""Get the multimodal config from the tokenizer filename."""
+    if isinstance(tokenizer_filename, Path):
+        tokenizer_filename = str(tokenizer_filename)
+
     _version_str = tokenizer_filename.split(".")[-1]
     if _version_str == "model" or "m" not in _version_str:
         return None
@@ -64,7 +70,7 @@ def get_mm_config(tokenizer_filename: str) -> Optional[MultimodalConfig]:
 class SentencePieceTokenizer(Tokenizer):
     r"""[SentencePiece](https://github.com/google/sentencepiece) tokenizer."""
 
-    def __init__(self, model_path: str, tokenizer_version: Optional[TokenizerVersion] = None) -> None:
+    def __init__(self, model_path: Union[str, Path], tokenizer_version: Optional[TokenizerVersion] = None) -> None:
         r"""Initialize the `SentencePieceTokenizer`.
 
         Args:
@@ -74,15 +80,23 @@ def __init__(self, model_path: str, tokenizer_version: Optional[TokenizerVersion
         self._logger = logging.getLogger(self.__class__.__name__)
         # reload tokenizer
         assert os.path.isfile(model_path), model_path
-        self._model = SentencePieceProcessor(model_file=model_path)
+        self._model = SentencePieceProcessor(
+            model_file=model_path if isinstance(model_path, str) else model_path.as_posix()
+        )
 
         assert self._model.vocab_size() == self._model.get_piece_size()
         self._vocab = [self._model.id_to_piece(i) for i in range(self.n_words)]
 
         self._version: TokenizerVersion = tokenizer_version or get_spm_version(model_path, raise_deprecated=False)
 
+        self._file_path = Path(model_path)
         super().__init__()
 
+    @property
+    def file_path(self) -> Path:
+        r"""The path to the tokenizer model."""
+        return self._file_path
+
     @property
     def version(self) -> TokenizerVersion:
         r"""The version of the tokenizer."""

src/mistral_common/tokens/tokenizers/tekken.py

@@ -137,7 +137,7 @@ def __init__(
         version: TokenizerVersion,
         *,
         name: str = "tekkenizer",
-        _path: Optional[str] = None,
+        _path: Optional[Union[str, Path]] = None,
         mm_config: Optional[MultimodalConfig] = None,
     ):
         r"""Initialize the tekken tokenizer.
@@ -199,6 +199,14 @@ def __init__(
         self._special_tokens_reverse_vocab = {t["token_str"]: t["rank"] for t in special_tokens}
         self._vocab = [self.id_to_piece(i) for i in range(vocab_size)]
         self._special_token_policy = SpecialTokenPolicy.IGNORE
+        self._file_path = Path(_path) if _path is not None else None
+
+    @property
+    def file_path(self) -> Path:
+        r"""The path to the tokenizer file."""
+        if self._file_path is None:
+            raise ValueError("The tokenizer was not loaded from a file.")
+        return self._file_path
 
     @classmethod
     def from_file(cls: Type["Tekkenizer"], path: Union[str, Path]) -> "Tekkenizer":
@@ -255,6 +263,7 @@ def from_file(cls: Type["Tekkenizer"], path: Union[str, Path]) -> "Tekkenizer":
             version=version,
             name=path.name.replace(".json", ""),
             mm_config=model_data.get("multimodal"),
+            _path=path,
         )
 
     @property

tests/test_mistral_tokenizer.py

@@ -1,4 +1,5 @@
-from typing import Optional, Union
+import multiprocessing
+from typing import List, Optional, Tuple, Union
 from unittest.mock import patch
 
 import pytest
@@ -103,3 +104,34 @@ def _mocked_hf_download(
 
             tokenizer = MistralTokenizer.from_hf_hub("mistralai/Pixtral-Large-Instruct-2411")
             assert isinstance(tokenizer.instruct_tokenizer, InstructTokenizerV3)
+
+
+def _worker_decode_function(tokenizer_instance_and_token_ids: Tuple[MistralTokenizer, List[int]]) -> str:
+    tokenizer_instance, token_ids = tokenizer_instance_and_token_ids
+    return tokenizer_instance.decode(token_ids)
+
+
+@pytest.mark.parametrize(
+    ["tokenizer_file", "token_ids", "expected"],
+    [
+        (
+            "tokenizer.model.v1",
+            [1, 733, 16289, 28793, 17121, 22526, 13, 13, 28708, 733, 28748, 16289, 28793],
+            "[INST] SYSTEM\n\na [/INST]",
+        ),
+        (
+            "tekken_240911.json",
+            [1091, 3174, 3074, 1093, 126205, 1267, 1097, 1766, 1047, 3174, 3074, 1093],
+            "[INST] SYSTEM\n\na [/INST]",
+        ),
+    ],
+)
+def test_tokenizer_is_pickleable_with_multiprocessing(tokenizer_file: str, token_ids: List[int], expected: str) -> None:
+    tokenizer_path = str(MistralTokenizer._data_path() / tokenizer_file)
+    tokenizer = MistralTokenizer.from_file(tokenizer_path)
+
+    with multiprocessing.Pool(processes=2) as pool:
+        results = pool.map(_worker_decode_function, [(tokenizer, token_ids)])
+
+    assert len(results) == 1
+    assert results[0] == expected