tokenization_herbert.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. # Copyright 2020 The Google AI Language Team Authors, Allegro.pl, Facebook Inc. and the HuggingFace Inc. team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
  15. from tokenizers.models import BPE
  16. from ...tokenization_utils_tokenizers import TokenizersBackend
  17. from ...utils import logging
  18. logger = logging.get_logger(__name__)
  19. VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt"}
  20. class HerbertTokenizer(TokenizersBackend):
  21. """
  22. Construct a BPE tokenizer for HerBERT (backed by HuggingFace's tokenizers library).
  23. Peculiarities:
  24. - uses BERT's pre-tokenizer: BertPreTokenizer splits tokens on spaces, and also on punctuation. Each occurrence of
  25. a punctuation character will be treated separately.
  26. This tokenizer inherits from [`TokenizersBackend`] which contains most of the methods. Users should refer to the
  27. superclass for more information regarding methods.
  28. Args:
  29. vocab_file (`str`):
  30. Path to the vocabulary file.
  31. merges_file (`str`):
  32. Path to the merges file.
  33. cls_token (`str`, *optional*, defaults to `"<s>"`):
  34. The classifier token.
  35. unk_token (`str`, *optional*, defaults to `"<unk>"`):
  36. The unknown token.
  37. pad_token (`str`, *optional*, defaults to `"<pad>"`):
  38. The padding token.
  39. mask_token (`str`, *optional*, defaults to `"<mask>"`):
  40. The mask token.
  41. sep_token (`str`, *optional*, defaults to `"</s>"`):
  42. The separator token.
  43. vocab (`str`, `dict` or `list`, *optional*):
  44. Custom vocabulary dictionary.
  45. merges (`str` or `list[str]`, *optional*):
  46. Custom merges list.
  47. """
  48. vocab_files_names = VOCAB_FILES_NAMES
  49. model_input_names = ["input_ids", "attention_mask"]
  50. model = BPE
  51. def __init__(
  52. self,
  53. vocab: str | dict[str, int] | None = None,
  54. merges: str | list[str] | None = None,
  55. cls_token: str = "<s>",
  56. unk_token: str = "<unk>",
  57. pad_token: str = "<pad>",
  58. mask_token: str = "<mask>",
  59. sep_token: str = "</s>",
  60. vocab_file: str | None = None,
  61. merges_file: str | None = None,
  62. **kwargs,
  63. ):
  64. self._vocab = vocab if vocab is not None else {str(unk_token): 0}
  65. self._merges = merges or []
  66. self._tokenizer = Tokenizer(
  67. BPE(
  68. vocab=self._vocab,
  69. merges=self._merges,
  70. dropout=None,
  71. unk_token=str(unk_token),
  72. end_of_word_suffix="</w>",
  73. )
  74. )
  75. self._tokenizer.normalizer = normalizers.BertNormalizer(
  76. lowercase=False, strip_accents=False, clean_text=True, handle_chinese_chars=True
  77. )
  78. self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
  79. self._tokenizer.decoder = decoders.BPEDecoder(suffix="</w>")
  80. super().__init__(
  81. cls_token=cls_token,
  82. unk_token=unk_token,
  83. pad_token=pad_token,
  84. mask_token=mask_token,
  85. sep_token=sep_token,
  86. **kwargs,
  87. )
  88. self._tokenizer.post_processor = processors.BertProcessing(
  89. sep=(self.sep_token, 2),
  90. cls=(self.cls_token, 0),
  91. )
  92. __all__ = ["HerbertTokenizer"]