(self, hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids)
| 376 | |
| 377 | |
| 378 | def process_entity_masks(self, hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids): |
| 379 | repeat_dim = hidden_states.shape[1] |
| 380 | max_masks = 0 |
| 381 | attention_mask = None |
| 382 | prompt_embs = [prompt_emb] |
| 383 | if entity_masks is not None: |
| 384 | # entity_masks |
| 385 | batch_size, max_masks = entity_masks.shape[0], entity_masks.shape[1] |
| 386 | entity_masks = entity_masks.repeat(1, 1, repeat_dim, 1, 1) |
| 387 | entity_masks = [entity_masks[:, i, None].squeeze(1) for i in range(max_masks)] |
| 388 | # global mask |
| 389 | global_mask = torch.ones_like(entity_masks[0]).to(device=hidden_states.device, dtype=hidden_states.dtype) |
| 390 | entity_masks = entity_masks + [global_mask] # append global to last |
| 391 | # attention mask |
| 392 | attention_mask = self.construct_mask(entity_masks, prompt_emb.shape[1], hidden_states.shape[1]) |
| 393 | attention_mask = attention_mask.to(device=hidden_states.device, dtype=hidden_states.dtype) |
| 394 | attention_mask = attention_mask.unsqueeze(1) |
| 395 | # embds: n_masks * b * seq * d |
| 396 | local_embs = [entity_prompt_emb[:, i, None].squeeze(1) for i in range(max_masks)] |
| 397 | prompt_embs = local_embs + prompt_embs # append global to last |
| 398 | prompt_embs = [self.context_embedder(prompt_emb) for prompt_emb in prompt_embs] |
| 399 | prompt_emb = torch.cat(prompt_embs, dim=1) |
| 400 | |
| 401 | # positional embedding |
| 402 | text_ids = torch.cat([text_ids] * (max_masks + 1), dim=1) |
| 403 | image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1)) |
| 404 | return prompt_emb, image_rotary_emb, attention_mask |
| 405 | |
| 406 | |
| 407 | def forward( |
no test coverage detected