Spaces:
Running
Running
| # ------------------------------------------------------------------------ | |
| # Grounding DINO | |
| # url: https://github.com/IDEA-Research/GroundingDINO | |
| # Copyright (c) 2023 IDEA. All Rights Reserved. | |
| # Licensed under the Apache License, Version 2.0 [see LICENSE for details] | |
| # ------------------------------------------------------------------------ | |
| import torch | |
| import torch.nn.functional as F | |
| import torch.utils.checkpoint as checkpoint | |
| from torch import Tensor, nn | |
| from torchvision.ops.boxes import nms | |
| from transformers import BertConfig, BertModel, BertPreTrainedModel | |
| from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions | |
| class BertModelWarper(nn.Module): | |
| def __init__(self, bert_model): | |
| super().__init__() | |
| # self.bert = bert_modelc | |
| self.config = bert_model.config | |
| self.embeddings = bert_model.embeddings | |
| self.encoder = bert_model.encoder | |
| self.pooler = bert_model.pooler | |
| self.get_extended_attention_mask = bert_model.get_extended_attention_mask | |
| self.invert_attention_mask = bert_model.invert_attention_mask | |
| self.get_head_mask = bert_model.get_head_mask | |
| def forward( | |
| self, | |
| input_ids=None, | |
| attention_mask=None, | |
| token_type_ids=None, | |
| position_ids=None, | |
| head_mask=None, | |
| inputs_embeds=None, | |
| encoder_hidden_states=None, | |
| encoder_attention_mask=None, | |
| past_key_values=None, | |
| use_cache=None, | |
| output_attentions=None, | |
| output_hidden_states=None, | |
| return_dict=None, | |
| ): | |
| r""" | |
| encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): | |
| Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if | |
| the model is configured as a decoder. | |
| encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): | |
| Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in | |
| the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``: | |
| - 1 for tokens that are **not masked**, | |
| - 0 for tokens that are **masked**. | |
| past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): | |
| Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. | |
| If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` | |
| (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` | |
| instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. | |
| use_cache (:obj:`bool`, `optional`): | |
| If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up | |
| decoding (see :obj:`past_key_values`). | |
| """ | |
| output_attentions = ( | |
| output_attentions if output_attentions is not None else self.config.output_attentions | |
| ) | |
| output_hidden_states = ( | |
| output_hidden_states | |
| if output_hidden_states is not None | |
| else self.config.output_hidden_states | |
| ) | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
| if self.config.is_decoder: | |
| use_cache = use_cache if use_cache is not None else self.config.use_cache | |
| else: | |
| use_cache = False | |
| if input_ids is not None and inputs_embeds is not None: | |
| raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") | |
| elif input_ids is not None: | |
| input_shape = input_ids.size() | |
| batch_size, seq_length = input_shape | |
| elif inputs_embeds is not None: | |
| input_shape = inputs_embeds.size()[:-1] | |
| batch_size, seq_length = input_shape | |
| else: | |
| raise ValueError("You have to specify either input_ids or inputs_embeds") | |
| device = input_ids.device if input_ids is not None else inputs_embeds.device | |
| # past_key_values_length | |
| past_key_values_length = ( | |
| past_key_values[0][0].shape[2] if past_key_values is not None else 0 | |
| ) | |
| if attention_mask is None: | |
| attention_mask = torch.ones( | |
| ((batch_size, seq_length + past_key_values_length)), device=device | |
| ) | |
| if token_type_ids is None: | |
| token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) | |
| # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] | |
| # ourselves in which case we just need to make it broadcastable to all heads. | |
| extended_attention_mask: torch.Tensor = self.get_extended_attention_mask( | |
| attention_mask, input_shape, device | |
| ) | |
| # If a 2D or 3D attention mask is provided for the cross-attention | |
| # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] | |
| if self.config.is_decoder and encoder_hidden_states is not None: | |
| encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() | |
| encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) | |
| if encoder_attention_mask is None: | |
| encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) | |
| encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) | |
| else: | |
| encoder_extended_attention_mask = None | |
| # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': | |
| # import ipdb; ipdb.set_trace() | |
| # Prepare head mask if needed | |
| # 1.0 in head_mask indicate we keep the head | |
| # attention_probs has shape bsz x n_heads x N x N | |
| # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] | |
| # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] | |
| head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) | |
| embedding_output = self.embeddings( | |
| input_ids=input_ids, | |
| position_ids=position_ids, | |
| token_type_ids=token_type_ids, | |
| inputs_embeds=inputs_embeds, | |
| past_key_values_length=past_key_values_length, | |
| ) | |
| encoder_outputs = self.encoder( | |
| embedding_output, | |
| attention_mask=extended_attention_mask, | |
| head_mask=head_mask, | |
| encoder_hidden_states=encoder_hidden_states, | |
| encoder_attention_mask=encoder_extended_attention_mask, | |
| past_key_values=past_key_values, | |
| use_cache=use_cache, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict, | |
| ) | |
| sequence_output = encoder_outputs[0] | |
| pooled_output = self.pooler(sequence_output) if self.pooler is not None else None | |
| if not return_dict: | |
| return (sequence_output, pooled_output) + encoder_outputs[1:] | |
| return BaseModelOutputWithPoolingAndCrossAttentions( | |
| last_hidden_state=sequence_output, | |
| pooler_output=pooled_output, | |
| past_key_values=encoder_outputs.past_key_values, | |
| hidden_states=encoder_outputs.hidden_states, | |
| attentions=encoder_outputs.attentions, | |
| cross_attentions=encoder_outputs.cross_attentions, | |
| ) | |
| class TextEncoderShell(nn.Module): | |
| def __init__(self, text_encoder): | |
| super().__init__() | |
| self.text_encoder = text_encoder | |
| self.config = self.text_encoder.config | |
| def forward(self, **kw): | |
| # feed into text encoder | |
| return self.text_encoder(**kw) | |
| def generate_masks_with_special_tokens(tokenized, special_tokens_list, tokenizer): | |
| """Generate attention mask between each pair of special tokens | |
| Args: | |
| input_ids (torch.Tensor): input ids. Shape: [bs, num_token] | |
| special_tokens_mask (list): special tokens mask. | |
| Returns: | |
| torch.Tensor: attention mask between each special tokens. | |
| """ | |
| input_ids = tokenized["input_ids"] | |
| bs, num_token = input_ids.shape | |
| # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens | |
| special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool() | |
| for special_token in special_tokens_list: | |
| special_tokens_mask |= input_ids == special_token | |
| # idxs: each row is a list of indices of special tokens | |
| idxs = torch.nonzero(special_tokens_mask) | |
| # generate attention mask and positional ids | |
| attention_mask = ( | |
| torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1) | |
| ) | |
| position_ids = torch.zeros((bs, num_token), device=input_ids.device) | |
| previous_col = 0 | |
| for i in range(idxs.shape[0]): | |
| row, col = idxs[i] | |
| if (col == 0) or (col == num_token - 1): | |
| attention_mask[row, col, col] = True | |
| position_ids[row, col] = 0 | |
| else: | |
| attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True | |
| position_ids[row, previous_col + 1 : col + 1] = torch.arange( | |
| 0, col - previous_col, device=input_ids.device | |
| ) | |
| previous_col = col | |
| # # padding mask | |
| # padding_mask = tokenized['attention_mask'] | |
| # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool() | |
| return attention_mask, position_ids.to(torch.long) | |
| def generate_masks_with_special_tokens_and_transfer_map(tokenized, special_tokens_list, tokenizer): | |
| """Generate attention mask between each pair of special tokens | |
| Args: | |
| input_ids (torch.Tensor): input ids. Shape: [bs, num_token] | |
| special_tokens_mask (list): special tokens mask. | |
| Returns: | |
| torch.Tensor: attention mask between each special tokens. | |
| """ | |
| input_ids = tokenized["input_ids"] | |
| bs, num_token = input_ids.shape | |
| # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens | |
| special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool() | |
| for special_token in special_tokens_list: | |
| special_tokens_mask |= input_ids == special_token | |
| # idxs: each row is a list of indices of special tokens | |
| idxs = torch.nonzero(special_tokens_mask) | |
| # generate attention mask and positional ids | |
| attention_mask = ( | |
| torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1) | |
| ) | |
| position_ids = torch.zeros((bs, num_token), device=input_ids.device) | |
| cate_to_token_mask_list = [[] for _ in range(bs)] | |
| previous_col = 0 | |
| for i in range(idxs.shape[0]): | |
| row, col = idxs[i] | |
| if (col == 0) or (col == num_token - 1): | |
| attention_mask[row, col, col] = True | |
| position_ids[row, col] = 0 | |
| else: | |
| attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True | |
| position_ids[row, previous_col + 1 : col + 1] = torch.arange( | |
| 0, col - previous_col, device=input_ids.device | |
| ) | |
| c2t_maski = torch.zeros((num_token), device=input_ids.device).bool() | |
| c2t_maski[previous_col + 1 : col] = True | |
| cate_to_token_mask_list[row].append(c2t_maski) | |
| previous_col = col | |
| cate_to_token_mask_list = [ | |
| torch.stack(cate_to_token_mask_listi, dim=0) | |
| for cate_to_token_mask_listi in cate_to_token_mask_list | |
| ] | |
| # # padding mask | |
| # padding_mask = tokenized['attention_mask'] | |
| # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool() | |
| return attention_mask, position_ids.to(torch.long), cate_to_token_mask_list | |