
"""simple docstring"""




import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ...utils import logging
from ..auto import CONFIG_MAPPING


A__ :       List[Any]				= logging.get_logger(__name__)

A__ :       Tuple				= {
    'Salesforce/instruct-blip-flan-t5': 'https://huggingface.co/Salesforce/instruct-blip-flan-t5/resolve/main/config.json',
}


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       UpperCamelCase_ =		'''instructblip_vision_model'''
       def __init__(  self			,					A_=1408			,					A_=6144			,					A_=39			,					A_=16			,					A_=224			,					A_=14			,					A_="gelu"			,					A_=1E-6			,					A_=0.0			,					A_=1E-10			,					A_=True			,					**A_			,					)  ->						Tuple:





         """simple docstring"""




         super().__init__(**A_							)

         _lowercase:      List[str]   			=  hidden_size
         _lowercase:      Tuple   			=  intermediate_size
         _lowercase:      List[str]   			=  num_hidden_layers
         _lowercase:      Union[str, Any]   			=  num_attention_heads
         _lowercase:      str   			=  patch_size
         _lowercase:      Union[str, Any]   			=  image_size
         _lowercase:      int   			=  initializer_range
         _lowercase:      List[str]   			=  attention_dropout
         _lowercase:      Tuple   			=  layer_norm_eps
         _lowercase:      Tuple   			=  hidden_act
         _lowercase:      Optional[int]   			=  qkv_bias



       @classmethod
       def 				lowercase_				(  cls			,					A_			,					**A_							)  ->						"PretrainedConfig":





         """simple docstring"""




         cls._set_token_in_kwargs(A_							)

         _lowercase			,					_lowercase:      str   			=  cls.get_config_dict(A_			,					**A_							)

         # get the vision config dict if we are loading from InstructBlipConfig
         if config_dict.get('''model_type'''							) == "instructblip":
           _lowercase:      Dict   			=  config_dict['''vision_config''']

         if "model_type" in config_dict and hasattr(cls			,					'''model_type'''							) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

         return cls.from_dict(A_			,					**A_							)


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       UpperCamelCase_ =		'''instructblip_qformer'''
       def __init__(  self			,					A_=3_0522			,					A_=768			,					A_=12			,					A_=12			,					A_=3072			,					A_="gelu"			,					A_=0.1			,					A_=0.1			,					A_=512			,					A_=0.02			,					A_=1E-12			,					A_=0			,					A_="absolute"			,					A_=2			,					A_=1408			,					**A_			,					)  ->						Optional[int]:





         """simple docstring"""




         super().__init__(pad_token_id=A_			,					**A_							)

         _lowercase:      Optional[Any]   			=  vocab_size
         _lowercase:      Optional[int]   			=  hidden_size
         _lowercase:      Dict   			=  num_hidden_layers
         _lowercase:      Union[str, Any]   			=  num_attention_heads
         _lowercase:      List[Any]   			=  hidden_act
         _lowercase:      Union[str, Any]   			=  intermediate_size
         _lowercase:      Union[str, Any]   			=  hidden_dropout_prob
         _lowercase:      Tuple   			=  attention_probs_dropout_prob
         _lowercase:      Optional[int]   			=  max_position_embeddings
         _lowercase:      Optional[int]   			=  initializer_range
         _lowercase:      Optional[Any]   			=  layer_norm_eps
         _lowercase:      Optional[Any]   			=  position_embedding_type
         _lowercase:      Tuple   			=  cross_attention_frequency
         _lowercase:      str   			=  encoder_hidden_size



       @classmethod
       def 				lowercase_				(  cls			,					A_			,					**A_							)  ->						"PretrainedConfig":





         """simple docstring"""




         cls._set_token_in_kwargs(A_							)

         _lowercase			,					_lowercase:      Tuple   			=  cls.get_config_dict(A_			,					**A_							)

         # get the qformer config dict if we are loading from InstructBlipConfig
         if config_dict.get('''model_type'''							) == "instructblip":
           _lowercase:      Tuple   			=  config_dict['''qformer_config''']

         if "model_type" in config_dict and hasattr(cls			,					'''model_type'''							) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

         return cls.from_dict(A_			,					**A_							)


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       UpperCamelCase_ =		'''instructblip'''
       UpperCamelCase_ =		True
       def __init__(  self			,					A_=None			,					A_=None			,					A_=None			,					A_=32			,					**A_							)  ->						str:





         """simple docstring"""




         super().__init__(**A_							)

         if vision_config is None:
           _lowercase:      str   			=  {}
           logger.info('''vision_config is None. initializing the InstructBlipVisionConfig with default values.'''							)

         if qformer_config is None:
           _lowercase:      List[str]   			=  {}
           logger.info('''qformer_config is None. Initializing the InstructBlipQFormerConfig with default values.'''							)

         if text_config is None:
           _lowercase:      Any   			=  {}
           logger.info('''text_config is None. Initializing the text config with default values (`OPTConfig`).'''							)

         _lowercase:      int   			=  InstructBlipVisionConfig(**A_							)
         _lowercase:      Optional[Any]   			=  InstructBlipQFormerConfig(**A_							)
         _lowercase:      Optional[Any]   			=  text_config['''model_type'''] if '''model_type''' in text_config else '''opt'''
         _lowercase:      Optional[Any]   			=  CONFIG_MAPPING[text_model_type](**A_							)

         _lowercase:      Tuple   			=  self.text_config.tie_word_embeddings
         _lowercase:      Dict   			=  self.text_config.is_encoder_decoder

         _lowercase:      List[Any]   			=  num_query_tokens
         _lowercase:      Any   			=  self.vision_config.hidden_size
         _lowercase:      List[Any]   			=  self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
         _lowercase:      Union[str, Any]   			=  1.0
         _lowercase:      Dict   			=  0.02
       @classmethod
       def 				lowercase_				(  cls			,					A_			,					A_			,					A_			,					**A_			,					)  ->						str:





         """simple docstring"""




         return cls(
             vision_config=vision_config.to_dict()			,					qformer_config=qformer_config.to_dict()			,					text_config=text_config.to_dict()			,					**A_			,					)



       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase:      List[str]   			=  copy.deepcopy(self.__dict__							)
         _lowercase:      Tuple   			=  self.vision_config.to_dict()
         _lowercase:      Tuple   			=  self.qformer_config.to_dict()
         _lowercase:      int   			=  self.text_config.to_dict()
         _lowercase:      Any   			=  self.__class__.model_type
         return output


"""simple docstring"""




import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       def __init__(  self			,					A_			,					A_							)  ->						List[str]:





         """simple docstring"""




         _lowercase:      List[str]   			=  params

         _lowercase:      str   			=  np.array(A_							)
         _lowercase:      Optional[int]   			=  np.array([len(A_							) for t in data]							)

         self.check()
         self.remove_long_sequences()
         self.remove_empty_sequences()
         self.remove_unknown_sequences()
         self.check()
         self.print_statistics()
       def __getitem__(  self			,					A_							)  ->						Optional[Any]:





         """simple docstring"""




         return (self.token_ids[index], self.lengths[index])
       def __len__(  self							)  ->						List[str]:





         """simple docstring"""




         return len(self.lengths							)
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         assert len(self.token_ids							) == len(self.lengths							)
         assert all(self.lengths[i] == len(self.token_ids[i]							) for i in range(len(self.lengths							)							)							)
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         _lowercase:      Tuple   			=  self.params.max_model_input_size
         _lowercase:      Tuple   			=  self.lengths > max_len
         logger.info(f'''Splitting {sum(A_							)} too long sequences.'''							)

         def divide_chunks(A_			,					A_							):
           return [l[i : i + n] for i in range(0			,					len(A_							)			,					A_							)]

         _lowercase:      Dict   			=  []
         _lowercase:      Union[str, Any]   			=  []
         if self.params.mlm:
           _lowercase			,					_lowercase:      int   			=  self.params.special_tok_ids['''cls_token'''], self.params.special_tok_ids['''sep_token''']
         else:
           _lowercase			,					_lowercase:      Dict   			=  self.params.special_tok_ids['''bos_token'''], self.params.special_tok_ids['''eos_token''']

         for seq_, len_ in zip(self.token_ids			,					self.lengths							):
           assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
           if len_ <= max_len:
             new_tok_ids.append(seq_							)
             new_lengths.append(len_							)
           else:
             _lowercase:      Optional[Any]   			=  []
             for sub_s in divide_chunks(seq_			,					max_len - 2							):
               if sub_s[0] != cls_id:
                 _lowercase:      Any   			=  np.insert(A_			,					0			,					A_							)
               if sub_s[-1] != sep_id:
                 _lowercase:      Optional[int]   			=  np.insert(A_			,					len(A_							)			,					A_							)
               assert len(A_							) <= max_len
               assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
               sub_seqs.append(A_							)

             new_tok_ids.extend(A_							)
             new_lengths.extend([len(A_							) for l in sub_seqs]							)

         _lowercase:      Optional[Any]   			=  np.array(A_							)
         _lowercase:      List[str]   			=  np.array(A_							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         _lowercase:      List[Any]   			=  len(self							)
         _lowercase:      Optional[Any]   			=  self.lengths > 11
         _lowercase:      int   			=  self.token_ids[indices]
         _lowercase:      List[str]   			=  self.lengths[indices]
         _lowercase:      Optional[int]   			=  len(self							)
         logger.info(f'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''							)
       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         if "unk_token" not in self.params.special_tok_ids:
           return
         else:
           _lowercase:      Dict   			=  self.params.special_tok_ids['''unk_token''']
         _lowercase:      int   			=  len(self							)
         _lowercase:      Tuple   			=  np.array([np.count_nonzero(a == unk_token_id							) for a in self.token_ids]							)
         _lowercase:      Dict   			=  (unk_occs / self.lengths) < 0.5
         _lowercase:      str   			=  self.token_ids[indices]
         _lowercase:      Dict   			=  self.lengths[indices]
         _lowercase:      int   			=  len(self							)
         logger.info(f'''Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'''							)
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         if not self.params.is_master:
           return
         logger.info(f'''{len(self							)} sequences'''							)
         # data_len = sum(self.lengths)
         # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
         # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

         # unk_idx = self.params.special_tok_ids['unk_token']
         # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
         # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')



       def 				lowercase_				(  self			,					A_							)  ->						Optional[int]:





         """simple docstring"""




         _lowercase:      List[str]   			=  [t[0] for t in batch]
         _lowercase:      Dict   			=  [t[1] for t in batch]
         assert len(A_							) == len(A_							)

         # Max for paddings
         _lowercase:      Tuple   			=  max(A_							)

         # Pad token ids
         if self.params.mlm:
           _lowercase:      str   			=  self.params.special_tok_ids['''pad_token''']
         else:
           _lowercase:      Union[str, Any]   			=  self.params.special_tok_ids['''unk_token''']
         _lowercase:      int   			=  [list(t.astype(A_							)							) + [pad_idx] * (max_seq_len_ - len(A_							)) for t in token_ids]
         assert len(tk_							) == len(A_							)
         assert all(len(A_							) == max_seq_len_ for t in tk_							)

         _lowercase:      str   			=  torch.tensor(tk_							)  # (bs, max_seq_len_)
         _lowercase:      int   			=  torch.tensor(A_							)  # (bs)
         return tk_t, lg_t
