tiktoken.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. from pathlib import Path
  2. from typing import Any
  3. from transformers.convert_slow_tokenizer import TikTokenConverter
  4. from transformers.tokenization_utils_fast import TIKTOKEN_VOCAB_FILE, TOKENIZER_FILE
  5. def convert_tiktoken_to_fast(encoding: Any, output_dir: str):
  6. """
  7. Converts given `tiktoken` encoding to `PretrainedTokenizerFast` and saves the configuration of converted tokenizer
  8. on disk.
  9. Args:
  10. encoding (`str` or `tiktoken.Encoding`):
  11. Tokenizer from `tiktoken` library. If `encoding` is `str`, the tokenizer will be loaded with
  12. `tiktoken.get_encoding(encoding)`.
  13. output_dir (`str`):
  14. Save path for converted tokenizer configuration file.
  15. """
  16. output_dir = Path(output_dir)
  17. output_dir.mkdir(exist_ok=True)
  18. save_file = output_dir / "tiktoken" / TIKTOKEN_VOCAB_FILE
  19. tokenizer_file = output_dir / TOKENIZER_FILE
  20. save_file_absolute = str(save_file.absolute())
  21. output_file_absolute = str(tokenizer_file.absolute())
  22. try:
  23. from tiktoken import get_encoding
  24. from tiktoken.load import dump_tiktoken_bpe
  25. if isinstance(encoding, str):
  26. encoding = get_encoding(encoding)
  27. dump_tiktoken_bpe(encoding._mergeable_ranks, save_file_absolute)
  28. except ImportError:
  29. raise ValueError("`tiktoken` is required to save a `tiktoken` file. Install it with `pip install tiktoken`.")
  30. tokenizer = TikTokenConverter(
  31. vocab_file=save_file_absolute, pattern=encoding._pat_str, additional_special_tokens=encoding._special_tokens
  32. ).converted()
  33. tokenizer.save(output_file_absolute)