tokenizer.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333
  1. # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team and Alibaba inc.
  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. """Tokenization classes."""
  15. from __future__ import absolute_import, division, print_function
  16. import collections
  17. import unicodedata
  18. import six
  19. def convert_to_unicode(text):
  20. """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
  21. if six.PY3:
  22. if isinstance(text, str):
  23. return text
  24. elif isinstance(text, bytes):
  25. return text.decode('utf-8', 'ignore')
  26. else:
  27. raise ValueError('Unsupported string type: %s' % (type(text)))
  28. elif six.PY2:
  29. if isinstance(text, str):
  30. return text.decode('utf-8', 'ignore')
  31. elif isinstance(text, unicode):
  32. return text
  33. else:
  34. raise ValueError('Unsupported string type: %s' % (type(text)))
  35. else:
  36. raise ValueError('Not running on Python2 or Python 3?')
  37. def printable_text(text):
  38. """Returns text encoded in a way suitable for print or `tf.logging`."""
  39. # These functions want `str` for both Python2 and Python3, but in one case
  40. # it's a Unicode string and in the other it's a byte string.
  41. if six.PY3:
  42. if isinstance(text, str):
  43. return text
  44. elif isinstance(text, bytes):
  45. return text.decode('utf-8', 'ignore')
  46. else:
  47. raise ValueError('Unsupported string type: %s' % (type(text)))
  48. elif six.PY2:
  49. if isinstance(text, str):
  50. return text
  51. elif isinstance(text, unicode):
  52. return text.encode('utf-8')
  53. else:
  54. raise ValueError('Unsupported string type: %s' % (type(text)))
  55. else:
  56. raise ValueError('Not running on Python2 or Python 3?')
  57. def load_vocab(vocab_file):
  58. """Loads a vocabulary file into a dictionary."""
  59. vocab = collections.OrderedDict()
  60. index = 0
  61. with open(vocab_file, 'r', encoding='utf-8') as reader:
  62. while True:
  63. token = convert_to_unicode(reader.readline())
  64. if not token:
  65. break
  66. token = token.strip()
  67. vocab[token] = index
  68. index += 1
  69. return vocab
  70. def convert_tokens_to_ids(vocab, tokens):
  71. """Converts a sequence of tokens into ids using the vocab."""
  72. ids = []
  73. for token in tokens:
  74. ids.append(vocab[token])
  75. return ids
  76. def whitespace_tokenize(text):
  77. """Runs basic whitespace cleaning and splitting on a piece of text."""
  78. text = text.strip()
  79. if not text:
  80. return []
  81. tokens = text.split()
  82. return tokens
  83. class FullTokenizer(object):
  84. """Runs end-to-end tokenziation."""
  85. def __init__(self, vocab_file, do_lower_case=True):
  86. self.vocab = load_vocab(vocab_file)
  87. self.inv_vocab = {v: k for k, v in self.vocab.items()}
  88. self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
  89. self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
  90. def tokenize(self, text):
  91. split_tokens = []
  92. for token in self.basic_tokenizer.tokenize(text):
  93. for sub_token in self.wordpiece_tokenizer.tokenize(token):
  94. split_tokens.append(sub_token)
  95. return split_tokens
  96. def convert_tokens_to_ids(self, tokens):
  97. return convert_tokens_to_ids(self.vocab, tokens)
  98. def convert_ids_to_tokens(self, ids):
  99. return [self.inv_vocab[i] for i in ids]
  100. class BasicTokenizer(object):
  101. """Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
  102. def __init__(self, do_lower_case=True):
  103. """Constructs a BasicTokenizer.
  104. Args:
  105. do_lower_case: Whether to lower case the input.
  106. """
  107. self.do_lower_case = do_lower_case
  108. def tokenize(self, text):
  109. """Tokenizes a piece of text."""
  110. text = convert_to_unicode(text)
  111. text = self._clean_text(text)
  112. # This was added on November 1st, 2018 for the multilingual and Chinese
  113. # models. This is also applied to the English models now, but it doesn't
  114. # matter since the English models were not trained on any Chinese data
  115. # and generally don't have any Chinese data in them (there are Chinese
  116. # characters in the vocabulary because Wikipedia does have some Chinese
  117. # words in the English Wikipedia.).
  118. text = self._tokenize_chinese_chars(text)
  119. orig_tokens = whitespace_tokenize(text)
  120. split_tokens = []
  121. for token in orig_tokens:
  122. if self.do_lower_case:
  123. token = token.lower()
  124. token = self._run_strip_accents(token)
  125. split_tokens.extend(self._run_split_on_punc(token))
  126. output_tokens = whitespace_tokenize(' '.join(split_tokens))
  127. return output_tokens
  128. def _run_strip_accents(self, text):
  129. """Strips accents from a piece of text."""
  130. text = unicodedata.normalize('NFD', text)
  131. output = []
  132. for char in text:
  133. cat = unicodedata.category(char)
  134. if cat == 'Mn':
  135. continue
  136. output.append(char)
  137. return ''.join(output)
  138. def _run_split_on_punc(self, text):
  139. """Splits punctuation on a piece of text."""
  140. chars = list(text)
  141. i = 0
  142. start_new_word = True
  143. output = []
  144. while i < len(chars):
  145. char = chars[i]
  146. if _is_punctuation(char):
  147. output.append([char])
  148. start_new_word = True
  149. else:
  150. if start_new_word:
  151. output.append([])
  152. start_new_word = False
  153. output[-1].append(char)
  154. i += 1
  155. return [''.join(x) for x in output]
  156. def _tokenize_chinese_chars(self, text):
  157. """Adds whitespace around any CJK character."""
  158. output = []
  159. for char in text:
  160. cp = ord(char)
  161. if self._is_chinese_char(cp):
  162. output.append(' ')
  163. output.append(char)
  164. output.append(' ')
  165. else:
  166. output.append(char)
  167. return ''.join(output)
  168. def _is_chinese_char(self, cp):
  169. """Checks whether CP is the codepoint of a CJK character."""
  170. # This defines a "chinese character" as anything in the CJK Unicode block:
  171. # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
  172. #
  173. # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
  174. # despite its name. The modern Korean Hangul alphabet is a different block,
  175. # as is Japanese Hiragana and Katakana. Those alphabets are used to write
  176. # space-separated words, so they are not treated specially and handled
  177. # like the all of the other languages.
  178. if ((cp >= 0x4E00 and cp <= 0x9FFF) or (cp >= 0x3400 and cp <= 0x4DBF)
  179. or (cp >= 0x20000 and cp <= 0x2A6DF)
  180. or (cp >= 0x2A700 and cp <= 0x2B73F)
  181. or (cp >= 0x2B740 and cp <= 0x2B81F)
  182. or (cp >= 0x2B820 and cp <= 0x2CEAF)
  183. or (cp >= 0xF900 and cp <= 0xFAFF)
  184. or (cp >= 0x2F800 and cp <= 0x2FA1F)):
  185. return True
  186. return False
  187. def _clean_text(self, text):
  188. """Performs invalid character removal and whitespace cleanup on text."""
  189. output = []
  190. for char in text:
  191. cp = ord(char)
  192. if cp == 0 or cp == 0xfffd or _is_control(char):
  193. continue
  194. if _is_whitespace(char):
  195. output.append(' ')
  196. else:
  197. output.append(char)
  198. return ''.join(output)
  199. class WordpieceTokenizer(object):
  200. """Runs WordPiece tokenization."""
  201. def __init__(self, vocab, unk_token='[UNK]', max_input_chars_per_word=100):
  202. self.vocab = vocab
  203. self.unk_token = unk_token
  204. self.max_input_chars_per_word = max_input_chars_per_word
  205. def tokenize(self, text):
  206. """Tokenizes a piece of text into its word pieces.
  207. This uses a greedy longest-match-first algorithm to perform tokenization
  208. using the given vocabulary.
  209. For example:
  210. >>> input = "unaffable"
  211. >>> output = ["un", "##aff", "##able"]
  212. Args:
  213. text: A single token or whitespace separated tokens. This should have
  214. already been passed through `BasicTokenizer.
  215. Returns:
  216. A list of wordpiece tokens.
  217. """
  218. text = convert_to_unicode(text)
  219. output_tokens = []
  220. for token in whitespace_tokenize(text):
  221. chars = list(token)
  222. if len(chars) > self.max_input_chars_per_word:
  223. output_tokens.append(self.unk_token)
  224. continue
  225. is_bad = False
  226. start = 0
  227. sub_tokens = []
  228. while start < len(chars):
  229. end = len(chars)
  230. cur_substr = None
  231. while start < end:
  232. substr = ''.join(chars[start:end])
  233. if start > 0:
  234. substr = '##' + substr
  235. if substr in self.vocab:
  236. cur_substr = substr
  237. break
  238. end -= 1
  239. if cur_substr is None:
  240. is_bad = True
  241. break
  242. sub_tokens.append(cur_substr)
  243. start = end
  244. if is_bad:
  245. output_tokens.append(self.unk_token)
  246. else:
  247. output_tokens.extend(sub_tokens)
  248. return output_tokens
  249. def _is_whitespace(char):
  250. """Checks whether `chars` is a whitespace character."""
  251. # \t, \n, and \r are technically control characters but we treat them
  252. # as whitespace since they are generally considered as such.
  253. if char == ' ' or char == '\t' or char == '\n' or char == '\r':
  254. return True
  255. cat = unicodedata.category(char)
  256. if cat == 'Zs':
  257. return True
  258. return False
  259. def _is_control(char):
  260. """Checks whether `chars` is a control character."""
  261. # These are technically control characters but we count them as whitespace
  262. # characters.
  263. if char == '\t' or char == '\n' or char == '\r':
  264. return False
  265. cat = unicodedata.category(char)
  266. if cat.startswith('C'):
  267. return True
  268. return False
  269. def _is_punctuation(char):
  270. """Checks whether `chars` is a punctuation character."""
  271. cp = ord(char)
  272. # We treat all non-letter/number ASCII as punctuation.
  273. # Characters such as "^", "$", and "`" are not in the Unicode
  274. # Punctuation class but we treat them as punctuation anyways, for
  275. # consistency.
  276. if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64)
  277. or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
  278. return True
  279. cat = unicodedata.category(char)
  280. if cat.startswith('P'):
  281. return True
  282. return False