| 931 | |
| 932 | |
| 933 | class GPT2BPETokenizer(Tokenizer): |
| 934 | def __init__(self, model_type_or_path, cache_dir=None, add_block_symbols=False, add_task_mask=False, |
| 935 | add_decoder_mask=False, **kwargs): |
| 936 | self.text_tokenizer = GPT2Tokenizer.from_pretrained(model_type_or_path, |
| 937 | cache_dir=cache_dir) |
| 938 | |
| 939 | # disable max len warnings by increasing max len |
| 940 | self.text_tokenizer.max_len = int(1e12) |
| 941 | self.num_tokens = len(self.text_tokenizer.encoder) |
| 942 | self.num_type_tokens = 2 |
| 943 | if model_type_or_path.startswith('roberta'): |
| 944 | self.num_command_tokens = 6 |
| 945 | self.num_text_tokens = self.num_tokens - 3 |
| 946 | self._command_tokens = [ |
| 947 | CommandToken('pad', '<|endoftext|>', self.text_tokenizer.encoder['</s>']), |
| 948 | CommandToken('eos', '<|endoftext|>', self.text_tokenizer.encoder['</s>']), |
| 949 | CommandToken('sep', '[SEP]', self.text_tokenizer.encoder['</s>']), |
| 950 | CommandToken('ENC', '[CLS]', self.text_tokenizer.encoder['<s>']), |
| 951 | CommandToken('MASK', '[MASK]', self.text_tokenizer.encoder['<mask>'], lstrip=True), |
| 952 | CommandToken('unk', '[UNK]', self.text_tokenizer.encoder['<unk>']) |
| 953 | ] |
| 954 | if add_block_symbols: |
| 955 | self._command_tokens.extend([ |
| 956 | CommandToken('sop', '<|startofpiece|>', self.num_tokens), |
| 957 | CommandToken('eop', '<|endofpiece|>', self.num_tokens + 1) |
| 958 | ]) |
| 959 | self.num_tokens += 2 |
| 960 | self.num_command_tokens += 2 |
| 961 | else: |
| 962 | self.num_command_tokens = 2 |
| 963 | self.num_text_tokens = self.num_tokens - 1 |
| 964 | self._command_tokens = [ |
| 965 | CommandToken('pad', '<|endoftext|>', self.text_tokenizer.encoder['<|endoftext|>']), |
| 966 | CommandToken('eos', '<|endoftext|>', self.text_tokenizer.encoder['<|endoftext|>']) |
| 967 | ] |
| 968 | if add_block_symbols: |
| 969 | self._command_tokens.extend([ |
| 970 | CommandToken('sop', '<|startofpiece|>', self.num_tokens), |
| 971 | CommandToken('eop', '<|endofpiece|>', self.num_tokens + 1), |
| 972 | CommandToken('ENC', '[CLS]', self.num_tokens + 2), |
| 973 | CommandToken('MASK', '[MASK]', self.num_tokens + 3, lstrip=True), |
| 974 | CommandToken('sep', '[SEP]', self.num_tokens + 4), |
| 975 | CommandToken('unk', '[UNK]', self.num_tokens + 5) |
| 976 | ]) |
| 977 | self.num_tokens += 6 |
| 978 | self.num_command_tokens += 6 |
| 979 | if add_block_symbols: |
| 980 | if add_task_mask: |
| 981 | self._command_tokens.extend([ |
| 982 | CommandToken('gMASK', '[gMASK]', self.num_tokens, lstrip=True), |
| 983 | CommandToken('sMASK', '[sMASK]', self.num_tokens + 1, lstrip=True) |
| 984 | ]) |
| 985 | self.num_tokens += 2 |
| 986 | self.num_command_tokens += 2 |
| 987 | if add_decoder_mask: |
| 988 | self._command_tokens.extend([ |
| 989 | CommandToken('dBLOCK', '[dBLOCK]', self.num_tokens) |
| 990 | ]) |