tokenization_splinter.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. # Copyright 2021 Tel AViv University, AllenAI and The HuggingFace Inc. team. All rights reserved.
  2. # All rights reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. """Tokenization classes for Splinter."""
  16. import collections
  17. from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
  18. from tokenizers.models import WordPiece
  19. from ...tokenization_utils_tokenizers import TokenizersBackend
  20. from ...utils import logging
  21. logger = logging.get_logger(__name__)
  22. VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
  23. def load_vocab(vocab_file):
  24. vocab = collections.OrderedDict()
  25. with open(vocab_file, "r", encoding="utf-8") as reader:
  26. tokens = reader.readlines()
  27. for index, token in enumerate(tokens):
  28. token = token.rstrip("\n")
  29. vocab[token] = index
  30. return vocab
  31. class SplinterTokenizer(TokenizersBackend):
  32. r"""
  33. Construct a Splinter tokenizer (backed by HuggingFace's tokenizers library). Based on WordPiece.
  34. This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should
  35. refer to this superclass for more information regarding those methods.
  36. Args:
  37. vocab_file (`str`, *optional*):
  38. Path to a vocabulary file.
  39. tokenizer_file (`str`, *optional*):
  40. Path to a tokenizers JSON file containing the serialization of a tokenizer.
  41. do_lower_case (`bool`, *optional*, defaults to `True`):
  42. Whether or not to lowercase the input when tokenizing.
  43. unk_token (`str`, *optional*, defaults to `"[UNK]"`):
  44. The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
  45. token instead.
  46. sep_token (`str`, *optional*, defaults to `"[SEP]"`):
  47. The separator token, which is used when building a sequence from multiple sequences.
  48. pad_token (`str`, *optional*, defaults to `"[PAD]"`):
  49. The token used for padding, for example when batching sequences of different lengths.
  50. cls_token (`str`, *optional*, defaults to `"[CLS]"`):
  51. The classifier token which is used when doing sequence classification.
  52. mask_token (`str`, *optional*, defaults to `"[MASK]"`):
  53. The token used for masking values.
  54. question_token (`str`, *optional*, defaults to `"[QUESTION]"`):
  55. The token used for constructing question representations.
  56. tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
  57. Whether or not to tokenize Chinese characters.
  58. strip_accents (`bool`, *optional*):
  59. Whether or not to strip all accents. If this option is not specified, then it will be determined by the
  60. value for `lowercase`.
  61. vocab (`str`, `dict` or `list`, *optional*):
  62. Custom vocabulary dictionary. If not provided, a minimal vocabulary is created.
  63. """
  64. vocab_files_names = VOCAB_FILES_NAMES
  65. model_input_names = ["input_ids", "attention_mask"]
  66. model = WordPiece
  67. def __init__(
  68. self,
  69. vocab: str | dict[str, int] | None = None,
  70. do_lower_case: bool = True,
  71. unk_token: str = "[UNK]",
  72. sep_token: str = "[SEP]",
  73. pad_token: str = "[PAD]",
  74. cls_token: str = "[CLS]",
  75. mask_token: str = "[MASK]",
  76. question_token: str = "[QUESTION]",
  77. tokenize_chinese_chars: bool = True,
  78. strip_accents: bool | None = None,
  79. **kwargs,
  80. ):
  81. self._vocab = (
  82. vocab
  83. if vocab is not None
  84. else {
  85. str(pad_token): 0,
  86. str(unk_token): 1,
  87. str(cls_token): 2,
  88. str(sep_token): 3,
  89. str(mask_token): 4,
  90. str(question_token): 5,
  91. ".": 6,
  92. }
  93. )
  94. self._tokenizer = Tokenizer(WordPiece(self._vocab, unk_token=str(unk_token)))
  95. self._tokenizer.normalizer = normalizers.BertNormalizer(
  96. clean_text=True,
  97. handle_chinese_chars=tokenize_chinese_chars,
  98. strip_accents=strip_accents,
  99. lowercase=do_lower_case,
  100. )
  101. self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
  102. self._tokenizer.decoder = decoders.WordPiece(prefix="##")
  103. super().__init__(
  104. unk_token=unk_token,
  105. sep_token=sep_token,
  106. pad_token=pad_token,
  107. cls_token=cls_token,
  108. mask_token=mask_token,
  109. question_token=question_token,
  110. do_lower_case=do_lower_case,
  111. tokenize_chinese_chars=tokenize_chinese_chars,
  112. strip_accents=strip_accents,
  113. **kwargs,
  114. )
  115. self.do_lower_case = do_lower_case
  116. self.tokenize_chinese_chars = tokenize_chinese_chars
  117. self.strip_accents = strip_accents
  118. self.question_token = question_token
  119. if self.question_token not in self.all_special_tokens:
  120. self.add_tokens([self.question_token], special_tokens=True)
  121. self.update_post_processor()
  122. @property
  123. def question_token_id(self):
  124. return self.convert_tokens_to_ids(self.question_token)
  125. def update_post_processor(self):
  126. cls = self.cls_token
  127. sep = self.sep_token
  128. question = self.question_token
  129. dot = "."
  130. cls_token_id = self.cls_token_id
  131. sep_token_id = self.sep_token_id
  132. question_token_id = self.question_token_id
  133. dot_token_id = self.convert_tokens_to_ids(".")
  134. if cls is None or sep is None:
  135. return
  136. if self.padding_side == "right":
  137. pair = f"{cls}:0 $A:0 {question} {dot} {sep}:0 $B:1 {sep}:1"
  138. else:
  139. pair = f"{cls}:0 $A:0 {sep}:0 $B:1 {question} {dot} {sep}:1"
  140. self._tokenizer.post_processor = processors.TemplateProcessing(
  141. single=f"{cls}:0 $A:0 {sep}:0",
  142. pair=pair,
  143. special_tokens=[
  144. (cls, cls_token_id),
  145. (sep, sep_token_id),
  146. (question, question_token_id),
  147. (dot, dot_token_id),
  148. ],
  149. )
  150. __all__ = ["SplinterTokenizer"]