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


_lowerCamelCase   :				Dict     			=						logging.get_logger(__name__)

_lowerCamelCase   :				Optional[Any]     			=						{
    '''salesforce/blip2-opt-2.7b''': '''https://huggingface.co/salesforce/blip2-opt-2.7b/resolve/main/config.json''',
}


class   lowercase      (       a       ):
							lowercase__  :    Dict									=							"""blip_2_vision_model"""




							def __init__(  self				:   Union[str, Any]		,							_UpperCamelCase				:   Any=1_408		,							_UpperCamelCase				:   int=6_144		,							_UpperCamelCase				:   Tuple=39		,							_UpperCamelCase				:   Optional[int]=16		,							_UpperCamelCase				:   List[Any]=224		,							_UpperCamelCase				:   List[Any]=14		,							_UpperCamelCase				:   Union[str, Any]="gelu"		,							_UpperCamelCase				:   Union[str, Any]=0.0_0_0_0_1		,							_UpperCamelCase				:   Optional[int]=0.0		,							_UpperCamelCase				:   List[str]=1e-10		,							_UpperCamelCase				:   List[Any]=True		,							**_UpperCamelCase				:   Dict		,							)   ->			Optional[Any]:



												'''simple docstring'''

												super().__init__(**_UpperCamelCase )

												SCREAMING_SNAKE_CASE           =    hidden_size
												SCREAMING_SNAKE_CASE           =    intermediate_size
												SCREAMING_SNAKE_CASE           =    num_hidden_layers
												SCREAMING_SNAKE_CASE           =    num_attention_heads
												SCREAMING_SNAKE_CASE           =    patch_size
												SCREAMING_SNAKE_CASE           =    image_size
												SCREAMING_SNAKE_CASE           =    initializer_range
												SCREAMING_SNAKE_CASE           =    attention_dropout
												SCREAMING_SNAKE_CASE           =    layer_norm_eps
												SCREAMING_SNAKE_CASE           =    hidden_act
												SCREAMING_SNAKE_CASE           =    qkv_bias





							@classmethod
							def 		__snake_case(  cls				:   Union[str, Any]		,							_UpperCamelCase				:   Union[str, os.PathLike]		,							**_UpperCamelCase				:   str )   ->			"PretrainedConfig":



												'''simple docstring'''

												cls._set_token_in_kwargs(_UpperCamelCase )

												SCREAMING_SNAKE_CASE     ,				SCREAMING_SNAKE_CASE           =    cls.get_config_dict(_UpperCamelCase		,							**_UpperCamelCase )

												# get the vision config dict if we are loading from Blip2Config
												if config_dict.get("model_type" ) == "blip-2":
																	SCREAMING_SNAKE_CASE           =    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(_UpperCamelCase		,							**_UpperCamelCase )



class   lowercase      (       a       ):
							lowercase__  :    Any									=							"""blip_2_qformer"""




							def __init__(  self				:   Tuple		,							_UpperCamelCase				:   Optional[Any]=30_522		,							_UpperCamelCase				:   str=768		,							_UpperCamelCase				:   List[str]=12		,							_UpperCamelCase				:   int=12		,							_UpperCamelCase				:   int=3_072		,							_UpperCamelCase				:   List[Any]="gelu"		,							_UpperCamelCase				:   Optional[int]=0.1		,							_UpperCamelCase				:   List[str]=0.1		,							_UpperCamelCase				:   Tuple=512		,							_UpperCamelCase				:   Tuple=0.0_2		,							_UpperCamelCase				:   Any=1e-12		,							_UpperCamelCase				:   Optional[int]=0		,							_UpperCamelCase				:   Union[str, Any]="absolute"		,							_UpperCamelCase				:   List[Any]=2		,							_UpperCamelCase				:   Optional[int]=1_408		,							**_UpperCamelCase				:   str		,							)   ->			Tuple:



												'''simple docstring'''

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

												SCREAMING_SNAKE_CASE           =    vocab_size
												SCREAMING_SNAKE_CASE           =    hidden_size
												SCREAMING_SNAKE_CASE           =    num_hidden_layers
												SCREAMING_SNAKE_CASE           =    num_attention_heads
												SCREAMING_SNAKE_CASE           =    hidden_act
												SCREAMING_SNAKE_CASE           =    intermediate_size
												SCREAMING_SNAKE_CASE           =    hidden_dropout_prob
												SCREAMING_SNAKE_CASE           =    attention_probs_dropout_prob
												SCREAMING_SNAKE_CASE           =    max_position_embeddings
												SCREAMING_SNAKE_CASE           =    initializer_range
												SCREAMING_SNAKE_CASE           =    layer_norm_eps
												SCREAMING_SNAKE_CASE           =    position_embedding_type
												SCREAMING_SNAKE_CASE           =    cross_attention_frequency
												SCREAMING_SNAKE_CASE           =    encoder_hidden_size





							@classmethod
							def 		__snake_case(  cls				:   Tuple		,							_UpperCamelCase				:   Union[str, os.PathLike]		,							**_UpperCamelCase				:   Any )   ->			"PretrainedConfig":



												'''simple docstring'''

												cls._set_token_in_kwargs(_UpperCamelCase )

												SCREAMING_SNAKE_CASE     ,				SCREAMING_SNAKE_CASE           =    cls.get_config_dict(_UpperCamelCase		,							**_UpperCamelCase )

												# get the qformer config dict if we are loading from Blip2Config
												if config_dict.get("model_type" ) == "blip-2":
																	SCREAMING_SNAKE_CASE           =    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(_UpperCamelCase		,							**_UpperCamelCase )



class   lowercase      (       a       ):
							lowercase__  :    Optional[int]									=							"""blip-2"""
							lowercase__  :    Union[str, Any]									=							True




							def __init__(  self				:   List[Any]		,							_UpperCamelCase				:   int=None		,							_UpperCamelCase				:   List[Any]=None		,							_UpperCamelCase				:   Union[str, Any]=None		,							_UpperCamelCase				:   int=32		,							**_UpperCamelCase				:   Optional[int] )   ->			List[Any]:



												'''simple docstring'''

												super().__init__(**_UpperCamelCase )

												if vision_config is None:
																	SCREAMING_SNAKE_CASE           =    {}
																	logger.info("vision_config is None. initializing the Blip2VisionConfig with default values." )

												if qformer_config is None:
																	SCREAMING_SNAKE_CASE           =    {}
																	logger.info("qformer_config is None. Initializing the Blip2QFormerConfig with default values." )

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

												SCREAMING_SNAKE_CASE           =    BlipaVisionConfig(**_UpperCamelCase )
												SCREAMING_SNAKE_CASE           =    BlipaQFormerConfig(**_UpperCamelCase )
												SCREAMING_SNAKE_CASE           =    text_config["model_type"] if "model_type" in text_config else "opt"
												SCREAMING_SNAKE_CASE           =    CONFIG_MAPPING[text_model_type](**_UpperCamelCase )

												SCREAMING_SNAKE_CASE           =    self.text_config.tie_word_embeddings
												SCREAMING_SNAKE_CASE           =    self.text_config.is_encoder_decoder

												SCREAMING_SNAKE_CASE           =    num_query_tokens
												SCREAMING_SNAKE_CASE           =    self.vision_config.hidden_size
												SCREAMING_SNAKE_CASE           =    self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
												SCREAMING_SNAKE_CASE           =    1.0
												SCREAMING_SNAKE_CASE           =    0.0_2




							@classmethod
							def 		__snake_case(  cls				:   Optional[Any]		,							_UpperCamelCase				:   BlipaVisionConfig		,							_UpperCamelCase				:   BlipaQFormerConfig		,							_UpperCamelCase				:   PretrainedConfig		,							**_UpperCamelCase				:   Optional[int]		,							)   ->			Dict:



												'''simple docstring'''

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





							def 		__snake_case(  self				:   List[str] )   ->			List[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =    copy.deepcopy(self.__dict__ )
												SCREAMING_SNAKE_CASE           =    self.vision_config.to_dict()
												SCREAMING_SNAKE_CASE           =    self.qformer_config.to_dict()
												SCREAMING_SNAKE_CASE           =    self.text_config.to_dict()
												SCREAMING_SNAKE_CASE           =    self.__class__.model_type
												return output





import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


_lowerCamelCase   :				Union[str, Any]     			=						logging.get_logger(__name__)


class   lowercase      :
							lowercase__  :    str									=							None





@experimental
def 		__lowerCamelCase							(UpperCAmelCase__       :							Any ,	UpperCAmelCase__       :							int ,	UpperCAmelCase__       :							Optional[int] ,	UpperCAmelCase__       :							List[Any] ,	UpperCAmelCase__       :							Union[str, Any] ,	UpperCAmelCase__       :							List[str] ,	UpperCAmelCase__       :							Any					):
					if ParallelBackendConfig.backend_name is None:
										return _map_with_multiprocessing_pool(
										    UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__					)

					return _map_with_joblib(UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__ ,	UpperCAmelCase__					)





def 		__lowerCamelCase							(UpperCAmelCase__       :							Optional[Any] ,	UpperCAmelCase__       :							List[Any] ,	UpperCAmelCase__       :							int ,	UpperCAmelCase__       :							Union[str, Any] ,	UpperCAmelCase__       :							Any ,	UpperCAmelCase__       :							int ,	UpperCAmelCase__       :							Optional[Any]					):
					SCREAMING_SNAKE_CASE           =    num_proc if num_proc <= len(UpperCAmelCase__					) else len(UpperCAmelCase__					)
					SCREAMING_SNAKE_CASE           =    []  # We organize the splits ourselve (contiguous splits)
					for index in range(UpperCAmelCase__					):
										SCREAMING_SNAKE_CASE           =    len(UpperCAmelCase__					) // num_proc
										SCREAMING_SNAKE_CASE           =    len(UpperCAmelCase__					) % num_proc
										SCREAMING_SNAKE_CASE           =    div * index + min(UpperCAmelCase__ ,	UpperCAmelCase__					)
										SCREAMING_SNAKE_CASE           =    start + div + (1 if index < mod else 0)
										split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)					)

					if len(UpperCAmelCase__					) != sum(len(i[1]					) for i in split_kwds					):
										raise ValueError(
										    F"Error dividing inputs iterable among processes. "
										    F"Total number of objects {len(UpperCAmelCase__					)}, "
										    F"length: {sum(len(i[1]					) for i in split_kwds					)}"					)

					logger.info(
					    F"Spawning {num_proc} processes for {len(UpperCAmelCase__					)} objects in slices of {[len(i[1]					) for i in split_kwds]}"					)
					SCREAMING_SNAKE_CASE     ,				SCREAMING_SNAKE_CASE           =    None, None
					if not disable_tqdm:
										SCREAMING_SNAKE_CASE     ,				SCREAMING_SNAKE_CASE           =    (RLock(),), tqdm.set_lock
					with Pool(UpperCAmelCase__ ,	initargs=UpperCAmelCase__ ,	initializer=UpperCAmelCase__					) as pool:
										SCREAMING_SNAKE_CASE           =    pool.map(UpperCAmelCase__ ,	UpperCAmelCase__					)
					logger.info(F"Finished {num_proc} processes"					)
					SCREAMING_SNAKE_CASE           =    [obj for proc_res in mapped for obj in proc_res]
					logger.info(F"Unpacked {len(UpperCAmelCase__					)} objects"					)

					return mapped





def 		__lowerCamelCase							(UpperCAmelCase__       :							Tuple ,	UpperCAmelCase__       :							Dict ,	UpperCAmelCase__       :							Any ,	UpperCAmelCase__       :							Optional[int] ,	UpperCAmelCase__       :							Tuple ,	UpperCAmelCase__       :							Tuple ,	UpperCAmelCase__       :							Any					):
					# progress bar is not yet supported for _map_with_joblib, because tqdm couldn't accurately be applied to joblib,
					# and it requires monkey-patching joblib internal classes which is subject to change
					import joblib

					with joblib.parallel_backend(ParallelBackendConfig.backend_name ,	n_jobs=UpperCAmelCase__					):
										return joblib.Parallel()(
										    joblib.delayed(UpperCAmelCase__					)((function, obj, types, None, True, None)					) for obj in iterable					)





@experimental
@contextlib.contextmanager
def 		__lowerCamelCase							(UpperCAmelCase__       :							str					):
					SCREAMING_SNAKE_CASE           =    backend_name

					if backend_name == "spark":
										from joblibspark import register_spark

										register_spark()

										# TODO: call create_cache_and_write_probe if "download" in steps
										# TODO: raise NotImplementedError when Dataset.map etc is called

					try:
										yield
					finally:
										SCREAMING_SNAKE_CASE           =    None




