| 151 | self.vision_end = tokenizer.encode('</vision>') |
| 152 | |
| 153 | def __call__(self, example, has_aux=False, add_bos_token=True, add_eos_token=True): |
| 154 | if has_aux: |
| 155 | example, *aux = example |
| 156 | else: |
| 157 | aux = tuple() |
| 158 | rand_state = random.Random(aux[-1]) # makes augmentations deterministic by line number |
| 159 | token_buffer = [] |
| 160 | loss_mask_buffer = [] |
| 161 | vision_mask = [] |
| 162 | |
| 163 | fields = example[self.config.fields_from_example] |
| 164 | if isinstance(fields, (tuple, list)): |
| 165 | if self.config.fields_index >= 0: |
| 166 | fields = fields[self.config.fields_index] |
| 167 | else: |
| 168 | # seed based on line number |
| 169 | fields = rand_state.choice(fields) |
| 170 | fields = fields.split(',') |
| 171 | |
| 172 | if add_bos_token and self.config.add_bos_token: |
| 173 | token_buffer.append(self.tokenizer.bos_token_id) |
| 174 | loss_mask_buffer.append(0.0) |
| 175 | vision_mask.append(False) |
| 176 | |
| 177 | for i, field in enumerate(fields): |
| 178 | if field.startswith('[') and field.endswith(']'): |
| 179 | # No loss for this field. |
| 180 | field = field[1:-1] |
| 181 | mask = 0.0 |
| 182 | else: |
| 183 | mask = 1.0 |
| 184 | |
| 185 | if field == '<|bos|>': |
| 186 | token_buffer.append(self.tokenizer.bos_token_id) |
| 187 | loss_mask_buffer.append(mask) |
| 188 | vision_mask.append(False) |
| 189 | elif field == '<|eos|>': |
| 190 | token_buffer.append(self.tokenizer.eos_token_id) |
| 191 | loss_mask_buffer.append(mask) |
| 192 | vision_mask.append(False) |
| 193 | elif 'vision' in field: |
| 194 | vision_tokens = example[field] |
| 195 | n_frames = int(len(vision_tokens) / self.config.n_tokens_per_frame) |
| 196 | if self.config.max_n_frames > 0 and n_frames > self.config.max_n_frames: # uniformly select |
| 197 | idxs = np.linspace(0, n_frames - 1, self.config.max_n_frames).astype(int) |
| 198 | new_vision_tokens = [] |
| 199 | for idx in idxs: |
| 200 | new_vision_tokens.extend(vision_tokens[idx * self.config.n_tokens_per_frame:(idx + 1) * self.config.n_tokens_per_frame]) |
| 201 | vision_tokens = new_vision_tokens |
| 202 | n_frames = self.config.max_n_frames |
| 203 | assert int(len(vision_tokens) / self.config.n_tokens_per_frame) == n_frames, (int(len(vision_tokens) / self.config.n_tokens_per_frame), n_frames) |
| 204 | |
| 205 | assert n_frames > 0, len(vision_tokens) |
| 206 | tokens = list(self.vision_start) |
| 207 | for j in range(n_frames): |
| 208 | tokens.extend(vision_tokens[j*self.config.n_tokens_per_frame:(j+1)*self.config.n_tokens_per_frame]) |
| 209 | if j == n_frames - 1: # last frame |
| 210 | tokens.append(self.config.eov_token) |