(self, entity_masks, prompt_seq_len, image_seq_len)
| 341 | |
| 342 | |
| 343 | def construct_mask(self, entity_masks, prompt_seq_len, image_seq_len): |
| 344 | N = len(entity_masks) |
| 345 | batch_size = entity_masks[0].shape[0] |
| 346 | total_seq_len = N * prompt_seq_len + image_seq_len |
| 347 | patched_masks = [self.patchify(entity_masks[i]) for i in range(N)] |
| 348 | attention_mask = torch.ones((batch_size, total_seq_len, total_seq_len), dtype=torch.bool).to(device=entity_masks[0].device) |
| 349 | |
| 350 | image_start = N * prompt_seq_len |
| 351 | image_end = N * prompt_seq_len + image_seq_len |
| 352 | # prompt-image mask |
| 353 | for i in range(N): |
| 354 | prompt_start = i * prompt_seq_len |
| 355 | prompt_end = (i + 1) * prompt_seq_len |
| 356 | image_mask = torch.sum(patched_masks[i], dim=-1) > 0 |
| 357 | image_mask = image_mask.unsqueeze(1).repeat(1, prompt_seq_len, 1) |
| 358 | # prompt update with image |
| 359 | attention_mask[:, prompt_start:prompt_end, image_start:image_end] = image_mask |
| 360 | # image update with prompt |
| 361 | attention_mask[:, image_start:image_end, prompt_start:prompt_end] = image_mask.transpose(1, 2) |
| 362 | # prompt-prompt mask |
| 363 | for i in range(N): |
| 364 | for j in range(N): |
| 365 | if i != j: |
| 366 | prompt_start_i = i * prompt_seq_len |
| 367 | prompt_end_i = (i + 1) * prompt_seq_len |
| 368 | prompt_start_j = j * prompt_seq_len |
| 369 | prompt_end_j = (j + 1) * prompt_seq_len |
| 370 | attention_mask[:, prompt_start_i:prompt_end_i, prompt_start_j:prompt_end_j] = False |
| 371 | |
| 372 | attention_mask = attention_mask.float() |
| 373 | attention_mask[attention_mask == 0] = float('-inf') |
| 374 | attention_mask[attention_mask == 1] = 0 |
| 375 | return attention_mask |
| 376 | |
| 377 | |
| 378 | def process_entity_masks(self, hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids): |
no test coverage detected