diff --git a/src/mistral_common/tokens/tokenizers/base.py b/src/mistral_common/tokens/tokenizers/base.py index 1ea1cb10..a31c1eed 100644 --- a/src/mistral_common/tokens/tokenizers/base.py +++ b/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) diff --git a/src/mistral_common/tokens/tokenizers/mistral.py b/src/mistral_common/tokens/tokenizers/mistral.py index 4a5cd9e5..db31cf09 100644 --- a/src/mistral_common/tokens/tokenizers/mistral.py +++ b/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,18 @@ 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, + self._chat_completion_request_validator._mode, + ) + @classmethod def _data_path(cls) -> Path: return Path(__file__).parents[2] / "data" @@ -238,7 +250,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. diff --git a/src/mistral_common/tokens/tokenizers/sentencepiece.py b/src/mistral_common/tokens/tokenizers/sentencepiece.py index 6ed65bbf..cfbf76e9 100644 --- a/src/mistral_common/tokens/tokenizers/sentencepiece.py +++ b/src/mistral_common/tokens/tokenizers/sentencepiece.py @@ -28,8 +28,10 @@ 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.""" + 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 +49,10 @@ 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.""" + 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 +68,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 +78,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.""" diff --git a/src/mistral_common/tokens/tokenizers/tekken.py b/src/mistral_common/tokens/tokenizers/tekken.py index e7f2dec8..cac3a6a6 100644 --- a/src/mistral_common/tokens/tokenizers/tekken.py +++ b/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 diff --git a/tests/test_mistral_tokenizer.py b/tests/test_mistral_tokenizer.py index 5d53ea6c..b5733f05 100644 --- a/tests/test_mistral_tokenizer.py +++ b/tests/test_mistral_tokenizer.py @@ -1,9 +1,11 @@ -from typing import Optional, Union +import multiprocessing +from typing import List, Optional, Tuple, Union from unittest.mock import patch import pytest from mistral_common.exceptions import TokenizerException +from mistral_common.protocol.instruct.validator import ValidationMode from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy from mistral_common.tokens.tokenizers.instruct import ( InstructTokenizerV1, @@ -103,3 +105,41 @@ 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_and_validation_mode: Tuple[MistralTokenizer, List[int], ValidationMode], +) -> str: + tokenizer_instance, token_ids, validation_mode = tokenizer_instance_and_token_ids_and_validation_mode + assert tokenizer_instance._chat_completion_request_validator._mode == validation_mode + return tokenizer_instance.decode(token_ids) + + +@pytest.mark.parametrize( + ["tokenizer_file", "validation_mode", "token_ids", "expected"], + [ + ( + "tokenizer.model.v1", + ValidationMode.test, + [1, 733, 16289, 28793, 17121, 22526, 13, 13, 28708, 733, 28748, 16289, 28793], + "[INST] SYSTEM\n\na [/INST]", + ), + ( + "tekken_240911.json", + ValidationMode.finetuning, + [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, validation_mode: ValidationMode, token_ids: List[int], expected: str +) -> None: + tokenizer_path = str(MistralTokenizer._data_path() / tokenizer_file) + tokenizer = MistralTokenizer.from_file(tokenizer_path, validation_mode) + + with multiprocessing.Pool(processes=2) as pool: + results = pool.map(_worker_decode_function, [(tokenizer, token_ids, validation_mode)]) + + assert len(results) == 1 + assert results[0] == expected