tokenization_openai.py 13.1 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.
thomwolf's avatar
thomwolf committed
3
4
5
6
7
8
9
10
11
12
13
14
15
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tokenization classes for OpenAI GPT."""
thomwolf's avatar
thomwolf committed
16
17
18
19
20
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import json
import logging
thomwolf's avatar
thomwolf committed
21
import os
thomwolf's avatar
thomwolf committed
22
import re
thomwolf's avatar
thomwolf committed
23
24
25
import sys
from io import open

thomwolf's avatar
thomwolf committed
26
from tqdm import tqdm
thomwolf's avatar
thomwolf committed
27
28

from .file_utils import cached_path
29
from .tokenization import BasicTokenizer
thomwolf's avatar
thomwolf committed
30
31
32
33
34
35
36
37
38
39
40
41
42
43

logger = logging.getLogger(__name__)

PRETRAINED_VOCAB_ARCHIVE_MAP = {
    'openai-gpt': "https://s3.amazonaws.com/models.huggingface.co/bert/openai-gpt-vocab.json",
}
PRETRAINED_MERGES_ARCHIVE_MAP = {
    'openai-gpt': "https://s3.amazonaws.com/models.huggingface.co/bert/openai-gpt-merges.txt",
}
PRETRAINED_VOCAB_POSITIONAL_EMBEDDINGS_SIZE_MAP = {
    'openai-gpt': 512,
}
VOCAB_NAME = 'vocab.json'
MERGES_NAME = 'merges.txt'
thomwolf's avatar
thomwolf committed
44
SPECIAL_TOKENS_NAME = 'special_tokens.txt'
thomwolf's avatar
thomwolf committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

def get_pairs(word):
    """
    Return set of symbol pairs in a word.
    word is represented as tuple of symbols (symbols being variable-length strings)
    """
    pairs = set()
    prev_char = word[0]
    for char in word[1:]:
        pairs.add((prev_char, char))
        prev_char = char
    return pairs

def text_standardize(text):
    """
    fixes some issues the spacy tokenizer had on books corpus
    also does some whitespace standardization
    """
    text = text.replace('—', '-')
    text = text.replace('–', '-')
    text = text.replace('―', '-')
    text = text.replace('…', '...')
    text = text.replace('´', "'")
    text = re.sub(r'''(-+|~+|!+|"+|;+|\?+|\++|,+|\)+|\(+|\\+|\/+|\*+|\[+|\]+|}+|{+|\|+|_+)''', r' \1 ', text)
    text = re.sub(r'\s*\n\s*', ' \n ', text)
    text = re.sub(r'[^\S\n]+', ' ', text)
    return text.strip()

thomwolf's avatar
thomwolf committed
73
class OpenAIGPTTokenizer(object):
thomwolf's avatar
thomwolf committed
74
    """
75
76
    BPE tokenizer. Peculiarities:
        - lower case all inputs
77
78
79
        - uses SpaCy tokenizer and ftfy for pre-BPE tokenization if they are installed, fallback to BERT's BasicTokenizer if not.
        - argument special_tokens and function set_special_tokens:
            can be used to add additional symbols (ex: "__classify__") to a vocabulary.
thomwolf's avatar
thomwolf committed
80
    """
thomwolf's avatar
thomwolf committed
81
    @classmethod
thomwolf's avatar
thomwolf committed
82
    def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, **kwargs):
thomwolf's avatar
thomwolf committed
83
84
85
86
        """
        Instantiate a PreTrainedBertModel from a pre-trained model file.
        Download and cache the pre-trained model file if needed.
        """
thomwolf's avatar
thomwolf committed
87
88
89
        if pretrained_model_name_or_path in PRETRAINED_VOCAB_ARCHIVE_MAP:
            vocab_file = PRETRAINED_VOCAB_ARCHIVE_MAP[pretrained_model_name_or_path]
            merges_file = PRETRAINED_MERGES_ARCHIVE_MAP[pretrained_model_name_or_path]
thomwolf's avatar
thomwolf committed
90
        else:
thomwolf's avatar
thomwolf committed
91
92
            vocab_file = os.path.join(pretrained_model_name_or_path, VOCAB_NAME)
            merges_file = os.path.join(pretrained_model_name_or_path, MERGES_NAME)
thomwolf's avatar
thomwolf committed
93
94
95
96
97
            special_tokens_file = os.path.join(pretrained_model_name_or_path, SPECIAL_TOKENS_NAME)
            if not os.path.exists(special_tokens_file):
                special_tokens_file = None
            else:
                logger.info("loading special tokens file {}".format(special_tokens_file))
thomwolf's avatar
thomwolf committed
98
99
100
101
        # redirect to the cache, if necessary
        try:
            resolved_vocab_file = cached_path(vocab_file, cache_dir=cache_dir)
            resolved_merges_file = cached_path(merges_file, cache_dir=cache_dir)
thomwolf's avatar
thomwolf committed
102
        except EnvironmentError:
thomwolf's avatar
thomwolf committed
103
104
            logger.error(
                "Model name '{}' was not found in model name list ({}). "
thomwolf's avatar
thomwolf committed
105
106
107
                "We assumed '{}' was a path or url but couldn't find files {} and {} "
                "at this path or url.".format(
                    pretrained_model_name_or_path,
thomwolf's avatar
thomwolf committed
108
                    ', '.join(PRETRAINED_VOCAB_ARCHIVE_MAP.keys()),
thomwolf's avatar
thomwolf committed
109
110
                    pretrained_model_name_or_path,
                    vocab_file, merges_file))
thomwolf's avatar
thomwolf committed
111
112
113
114
115
116
117
118
119
            return None
        if resolved_vocab_file == vocab_file and resolved_merges_file == merges_file:
            logger.info("loading vocabulary file {}".format(vocab_file))
            logger.info("loading merges file {}".format(merges_file))
        else:
            logger.info("loading vocabulary file {} from cache at {}".format(
                vocab_file, resolved_vocab_file))
            logger.info("loading merges file {} from cache at {}".format(
                merges_file, resolved_merges_file))
thomwolf's avatar
thomwolf committed
120
        if pretrained_model_name_or_path in PRETRAINED_VOCAB_POSITIONAL_EMBEDDINGS_SIZE_MAP:
thomwolf's avatar
thomwolf committed
121
122
            # if we're using a pretrained model, ensure the tokenizer wont index sequences longer
            # than the number of positional embeddings
thomwolf's avatar
thomwolf committed
123
            max_len = PRETRAINED_VOCAB_POSITIONAL_EMBEDDINGS_SIZE_MAP[pretrained_model_name_or_path]
thomwolf's avatar
thomwolf committed
124
125
            kwargs['max_len'] = min(kwargs.get('max_len', int(1e12)), max_len)
        # Instantiate tokenizer.
thomwolf's avatar
thomwolf committed
126
127
128
129
130
        if special_tokens_file and 'special_tokens' not in kwargs:
            special_tokens = open(special_tokens_file, encoding='utf-8').read().split('\n')[:-1]
        else:
            special_tokens = kwargs.pop('special_tokens', [])
        tokenizer = cls(resolved_vocab_file, resolved_merges_file, special_tokens=special_tokens, *inputs, **kwargs)
thomwolf's avatar
thomwolf committed
131
132
        return tokenizer

thomwolf's avatar
thomwolf committed
133
    def __init__(self, vocab_file, merges_file, special_tokens=None, max_len=None):
thomwolf's avatar
thomwolf committed
134
135
136
        try:
            import ftfy
            import spacy
137
138
            self.nlp = spacy.load('en', disable=['parser', 'tagger', 'ner', 'textcat'])
            self.fix_text = ftfy.fix_text
thomwolf's avatar
thomwolf committed
139
        except ImportError:
140
141
142
143
            logger.warning("ftfy or spacy is not installed using BERT BasicTokenizer instead of SpaCy & ftfy.")
            self.nlp = BasicTokenizer(do_lower_case=True,
                                      never_split=special_tokens if special_tokens is not None else [])
            self.fix_text = None
thomwolf's avatar
thomwolf committed
144

thomwolf's avatar
thomwolf committed
145
        self.max_len = max_len if max_len is not None else int(1e12)
thomwolf's avatar
thomwolf committed
146
        self.encoder = json.load(open(vocab_file, encoding="utf-8"))
thomwolf's avatar
thomwolf committed
147
        self.decoder = {v:k for k,v in self.encoder.items()}
thomwolf's avatar
thomwolf committed
148
        merges = open(merges_file, encoding='utf-8').read().split('\n')[1:-1]
thomwolf's avatar
thomwolf committed
149
150
151
        merges = [tuple(merge.split()) for merge in merges]
        self.bpe_ranks = dict(zip(merges, range(len(merges))))
        self.cache = {}
thomwolf's avatar
logging  
thomwolf committed
152
        self.set_special_tokens(special_tokens)
thomwolf's avatar
thomwolf committed
153

154
155
156
    def __len__(self):
        return len(self.encoder) + len(self.special_tokens)

thomwolf's avatar
thomwolf committed
157
    def set_special_tokens(self, special_tokens):
thomwolf's avatar
thomwolf committed
158
159
160
161
        """ Add a list of additional tokens to the encoder.
            The additional tokens are indexed starting from the last index of the
            current vocabulary in the order of the `special_tokens` list.
        """
thomwolf's avatar
logging  
thomwolf committed
162
163
164
165
        if not special_tokens:
            self.special_tokens = {}
            self.special_tokens_decoder = {}
            return
thomwolf's avatar
thomwolf committed
166
        self.special_tokens = dict((tok, len(self.encoder) + i) for i, tok in enumerate(special_tokens))
thomwolf's avatar
logging  
thomwolf committed
167
        self.special_tokens_decoder = {v:k for k, v in self.special_tokens.items()}
168
169
170
        if self.fix_text is None:
            # Using BERT's BasicTokenizer: we can update the tokenizer
            self.nlp.never_split = special_tokens
thomwolf's avatar
logging  
thomwolf committed
171
        logger.info("Special tokens {}".format(self.special_tokens))
thomwolf's avatar
thomwolf committed
172
173

    def bpe(self, token):
174
        word = tuple(token[:-1]) + (token[-1] + '</w>',)
thomwolf's avatar
thomwolf committed
175
176
177
178
179
180
181
182
        if token in self.cache:
            return self.cache[token]
        pairs = get_pairs(word)

        if not pairs:
            return token+'</w>'

        while True:
183
            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf')))
thomwolf's avatar
thomwolf committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
            if bigram not in self.bpe_ranks:
                break
            first, second = bigram
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                    new_word.extend(word[i:j])
                    i = j
                except:
                    new_word.extend(word[i:])
                    break

                if word[i] == first and i < len(word)-1 and word[i+1] == second:
                    new_word.append(first+second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break
            else:
                pairs = get_pairs(word)
        word = ' '.join(word)
        if word == '\n  </w>':
            word = '\n</w>'
        self.cache[token] = word
        return word

thomwolf's avatar
thomwolf committed
216
    def tokenize(self, text):
thomwolf's avatar
thomwolf committed
217
        """ Tokenize a string. """
thomwolf's avatar
thomwolf committed
218
        split_tokens = []
219
220
221
222
223
224
225
226
227
228
        if self.fix_text is None:
            # Using BERT's BasicTokenizer
            text = self.nlp.tokenize(text)
            for token in text:
                split_tokens.extend([t for t in self.bpe(token).split(' ')])
        else:
            # Using SpaCy & ftfy (original tokenization process of OpenAI GPT)
            text = self.nlp(text_standardize(self.fix_text(text)))
            for token in text:
                split_tokens.extend([t for t in self.bpe(token.text.lower()).split(' ')])
thomwolf's avatar
thomwolf committed
229
230
231
        return split_tokens

    def convert_tokens_to_ids(self, tokens):
thomwolf's avatar
thomwolf committed
232
        """ Converts a sequence of tokens into ids using the vocab. """
thomwolf's avatar
thomwolf committed
233
        ids = []
thomwolf's avatar
thomwolf committed
234
        if isinstance(tokens, str) or (sys.version_info[0] == 2 and isinstance(tokens, unicode)):
thomwolf's avatar
logging  
thomwolf committed
235
236
237
238
            if tokens in self.special_tokens:
                return self.special_tokens[tokens]
            else:
                return self.encoder.get(tokens, 0)
thomwolf's avatar
thomwolf committed
239
240
241
242
243
244
        for token in tokens:
            if token in self.special_tokens:
                ids.append(self.special_tokens[token])
            else:
                ids.append(self.encoder.get(token, 0))
        if len(ids) > self.max_len:
245
            logger.warning(
thomwolf's avatar
thomwolf committed
246
                "Token indices sequence length is longer than the specified maximum "
247
248
                " sequence length for this OpenAI GPT model ({} > {}). Running this"
                " sequence through the model will result in indexing errors".format(len(ids), self.max_len)
thomwolf's avatar
thomwolf committed
249
250
251
            )
        return ids

thomwolf's avatar
thomwolf committed
252
    def convert_ids_to_tokens(self, ids, skip_special_tokens=False):
thomwolf's avatar
thomwolf committed
253
254
255
        """Converts a sequence of ids in BPE tokens using the vocab."""
        tokens = []
        for i in ids:
thomwolf's avatar
logging  
thomwolf committed
256
            if i in self.special_tokens_decoder:
thomwolf's avatar
thomwolf committed
257
258
                if not skip_special_tokens:
                    tokens.append(self.special_tokens_decoder[i])
thomwolf's avatar
logging  
thomwolf committed
259
260
            else:
                tokens.append(self.decoder[i])
thomwolf's avatar
thomwolf committed
261
262
        return tokens

thomwolf's avatar
thomwolf committed
263
    def decode(self, ids, skip_special_tokens=False, clean_up_tokenization_spaces=False):
thomwolf's avatar
thomwolf committed
264
        """Converts a sequence of ids in a string."""
thomwolf's avatar
thomwolf committed
265
        tokens = self.convert_ids_to_tokens(ids, skip_special_tokens=skip_special_tokens)
thomwolf's avatar
thomwolf committed
266
        out_string = ''.join(tokens).replace('</w>', ' ').strip()
thomwolf's avatar
thomwolf committed
267
268
269
270
271
272
        if clean_up_tokenization_spaces:
            out_string = out_string.replace('<unk>', '')
            out_string = out_string.replace(' .', '.').replace(' ?', '?').replace(' !', '!').replace(' ,', ',').replace(' ,', ','
                    ).replace(" n't", "n't").replace(" 'm", "'m").replace(" 're", "'re").replace(" do not", " don't"
                    ).replace(" 's", "'s").replace(" t ", "'t ").replace(" s ", "'s ").replace(" m ", "'m "
                    ).replace(" 've", "'ve")
thomwolf's avatar
thomwolf committed
273
        return out_string
274
275

    def save_vocabulary(self, vocab_path):
276
277
278
279
        """Save the tokenizer vocabulary and merge files to a directory."""
        if not os.path.isdir(vocab_path):
            logger.error("Vocabulary path ({}) should be a directory".format(vocab_path))
            return
280
281
        vocab_file = os.path.join(vocab_path, VOCAB_NAME)
        merge_file = os.path.join(vocab_path, MERGES_NAME)
thomwolf's avatar
thomwolf committed
282
283
284
285
286
        special_tokens_file = os.path.join(vocab_path, SPECIAL_TOKENS_NAME)

        with open(vocab_file, 'w', encoding='utf-8') as f:
            f.write(json.dumps(self.encoder, ensure_ascii=False))

287
288
289
290
291
292
293
294
        index = 0
        with open(merge_file, "w", encoding="utf-8") as writer:
            writer.write(u'#version: 0.2\n')
            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
                if index != token_index:
                    logger.warning("Saving vocabulary to {}: BPE merge indices are not consecutive."
                                   " Please check that the tokenizer is not corrupted!".format(merge_file))
                    index = token_index
thomwolf's avatar
thomwolf committed
295
                writer.write(' '.join(bpe_tokens) + u'\n')
296
                index += 1
thomwolf's avatar
thomwolf committed
297
298
299
300
301
302

        with open(special_tokens_file, 'w', encoding='utf-8') as writer:
            for token in sorted(self.special_tokens.keys(), key=lambda kv: kv[1]):
                writer.write(token + u'\n')

        return vocab_file, merge_file, special_tokens_file