| 101 | |
| 102 | |
| 103 | def process_caption_bert(tokenizer, tokens, train=True, mask_rate=0.2, size_augment=True): |
| 104 | |
| 105 | output_tokens = [] |
| 106 | deleted_idx = [] |
| 107 | |
| 108 | for i, token in enumerate(tokens): |
| 109 | |
| 110 | # the sentence is tokenized twice |
| 111 | # text -> basic token (basic_tokenizer.tokenize) -> sub_token (wordpiece_tokenizer.tokenize) |
| 112 | sub_tokens = tokenizer.wordpiece_tokenizer.tokenize(token) |
| 113 | |
| 114 | prob = random.random() |
| 115 | |
| 116 | # first, 20% probability use the augmenation operations |
| 117 | if size_augment and prob < mask_rate and train: # mask/remove the tokens only during training |
| 118 | prob /= mask_rate |
| 119 | |
| 120 | # 50% change token to mask token |
| 121 | if prob < 0.5: |
| 122 | for sub_token in sub_tokens: |
| 123 | output_tokens.append("[MASK]") |
| 124 | # 10% randomly change token to random token from the BERT-vocab |
| 125 | elif prob < 0.6: |
| 126 | for sub_token in sub_tokens: |
| 127 | output_tokens.append(random.choice(list(tokenizer.vocab.keys()))) |
| 128 | |
| 129 | # -> 40% delete the token |
| 130 | else: |
| 131 | for sub_token in sub_tokens: |
| 132 | output_tokens.append(sub_token) |
| 133 | # record the index of sub_token |
| 134 | deleted_idx.append(len(output_tokens) - 1) |
| 135 | |
| 136 | # 80% probability keep the token |
| 137 | else: |
| 138 | for sub_token in sub_tokens: |
| 139 | # no masking token (will be ignored by loss function later) |
| 140 | output_tokens.append(sub_token) |
| 141 | |
| 142 | if len(deleted_idx) != 0: |
| 143 | output_tokens = [output_tokens[i] for i in range(len(output_tokens)) if i not in deleted_idx] |
| 144 | |
| 145 | # and first and last notations for BERT model |
| 146 | output_tokens = ['[CLS]'] + output_tokens + ['[SEP]'] |
| 147 | |
| 148 | # Convert each token to vocabulary indices |
| 149 | # [PAD] -> 0 |
| 150 | # [UNK] -> 100 |
| 151 | # [CLS] -> 101 |
| 152 | # [SEP] -> 102 |
| 153 | # [MASK] -> 103 |
| 154 | target = tokenizer.convert_tokens_to_ids(output_tokens) |
| 155 | |
| 156 | # convert to the torch.Tensor, torch.int64 (long) |
| 157 | target = torch.tensor(target) |
| 158 | |
| 159 | return target |
| 160 | |