


import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.modeling_flax_pytorch_utils import (
		    convert_pytorch_state_dict_to_flax,
		    load_flax_weights_in_pytorch_model,
		)
		from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
		import torch





class       _a							:




	def __init__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Dict	,      _SCREAMING_SNAKE_CASE		:	Tuple=14	,      _SCREAMING_SNAKE_CASE		:	str=7	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=True	,      _SCREAMING_SNAKE_CASE		:	Tuple=True	,      _SCREAMING_SNAKE_CASE		:	List[str]=False	,      _SCREAMING_SNAKE_CASE		:	int=True	,      _SCREAMING_SNAKE_CASE		:	List[Any]=99	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=32	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=5	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=4	,      _SCREAMING_SNAKE_CASE		:	List[str]=37	,      _SCREAMING_SNAKE_CASE		:	Tuple="gelu"	,      _SCREAMING_SNAKE_CASE		:	Tuple=0.1	,      _SCREAMING_SNAKE_CASE		:	List[Any]=0.1	,      _SCREAMING_SNAKE_CASE		:	List[str]=512	,      _SCREAMING_SNAKE_CASE		:	Any=0.02	,      )-> Union[str, Any]:
			lowerCAmelCase__     :    int					  =      parent
			lowerCAmelCase__     :    str					  =      batch_size
			lowerCAmelCase__     :    Tuple					  =      seq_length
			lowerCAmelCase__     :    int					  =      is_training
			lowerCAmelCase__     :    Tuple					  =      use_input_mask
			lowerCAmelCase__     :    Any					  =      use_token_type_ids
			lowerCAmelCase__     :    List[str]					  =      use_labels
			lowerCAmelCase__     :    int					  =      vocab_size
			lowerCAmelCase__     :    Any					  =      hidden_size
			lowerCAmelCase__     :    Dict					  =      num_hidden_layers
			lowerCAmelCase__     :    str					  =      num_attention_heads
			lowerCAmelCase__     :    Optional[int]					  =      intermediate_size
			lowerCAmelCase__     :    str					  =      hidden_act
			lowerCAmelCase__     :    List[Any]					  =      hidden_dropout_prob
			lowerCAmelCase__     :    Union[str, Any]					  =      attention_probs_dropout_prob
			lowerCAmelCase__     :    Tuple					  =      max_position_embeddings
			lowerCAmelCase__     :    Union[str, Any]					  =      initializer_range
			lowerCAmelCase__     :    Any					  =      None
			lowerCAmelCase__     :    Union[str, Any]					  =      vocab_size - 1
			lowerCAmelCase__     :    int					  =      vocab_size - 1
			lowerCAmelCase__     :    List[str]					  =      vocab_size - 1




	def 							UpperCAmelCase__(      self		:	Tuple	)-> Tuple:
			lowerCAmelCase__     :    str					  =      ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size	)

			lowerCAmelCase__     :    List[Any]					  =      None
			if self.use_input_mask:
					lowerCAmelCase__     :    str					  =      random_attention_mask([self.batch_size, self.seq_length]	)

			lowerCAmelCase__     :    int					  =      GPTaConfig(
			    vocab_size=self.vocab_size	,      n_embd=self.hidden_size	,      n_layer=self.num_hidden_layers	,      n_head=self.num_attention_heads	,      n_positions=self.max_position_embeddings	,      use_cache=_SCREAMING_SNAKE_CASE	,      bos_token_id=self.bos_token_id	,      eos_token_id=self.eos_token_id	,      pad_token_id=self.pad_token_id	,      )

			return (config, input_ids, input_mask)




	def 							UpperCAmelCase__(      self		:	Optional[int]	)-> str:
			lowerCAmelCase__     :    Dict					  =      self.prepare_config_and_inputs()
			lowerCAmelCase__  ,			lowerCAmelCase__  ,			lowerCAmelCase__     :    List[str]					  =      config_and_inputs
			lowerCAmelCase__     :    Any					  =      {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
			return config, inputs_dict




	def 							UpperCAmelCase__(      self		:	List[str]	)-> Tuple:
			lowerCAmelCase__  ,			lowerCAmelCase__  ,			lowerCAmelCase__     :    Optional[Any]					  =      self.prepare_config_and_inputs()

			lowerCAmelCase__     :    Union[str, Any]					  =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
			lowerCAmelCase__     :    Any					  =      ids_tensor([self.batch_size, self.seq_length]	,      vocab_size=2	)

			return (
			    config,
			    input_ids,
			    attention_mask,
			    encoder_hidden_states,
			    encoder_attention_mask,
			)




	def 							UpperCAmelCase__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	Dict	,      _SCREAMING_SNAKE_CASE		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	List[str]	)-> Any:
			lowerCAmelCase__     :    int					  =      20
			lowerCAmelCase__     :    Dict					  =      model_class_name(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Union[str, Any]					  =      model.init_cache(input_ids.shape[0]	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      jnp.ones((input_ids.shape[0], max_decoder_length)	,      dtype='''i4'''	)

			lowerCAmelCase__     :    Any					  =      jnp.broadcast_to(
			    jnp.arange(input_ids.shape[-1] - 1	)[None, :]	,      (input_ids.shape[0], input_ids.shape[-1] - 1)	)
			lowerCAmelCase__     :    Optional[int]					  =      model(
			    input_ids[:, :-1]	,      attention_mask=_SCREAMING_SNAKE_CASE	,      past_key_values=_SCREAMING_SNAKE_CASE	,      position_ids=_SCREAMING_SNAKE_CASE	,      )

			lowerCAmelCase__     :    str					  =      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,      dtype='''i4'''	)
			lowerCAmelCase__     :    int					  =      model(
			    input_ids[:, -1:]	,      attention_mask=_SCREAMING_SNAKE_CASE	,      past_key_values=outputs_cache.past_key_values	,      position_ids=_SCREAMING_SNAKE_CASE	,      )

			lowerCAmelCase__     :    Any					  =      model(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Optional[int]					  =      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
			self.parent.assertTrue(diff < 1E-3	,      msg=F'Max diff is {diff}'	)




	def 							UpperCAmelCase__(      self		:	Dict	,      _SCREAMING_SNAKE_CASE		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	int	)-> Any:
			lowerCAmelCase__     :    int					  =      20
			lowerCAmelCase__     :    int					  =      model_class_name(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    List[Any]					  =      jnp.concatenate(
			    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)]	,      axis=-1	,      )

			lowerCAmelCase__     :    int					  =      model.init_cache(input_ids.shape[0]	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    int					  =      jnp.broadcast_to(
			    jnp.arange(input_ids.shape[-1] - 1	)[None, :]	,      (input_ids.shape[0], input_ids.shape[-1] - 1)	)

			lowerCAmelCase__     :    List[str]					  =      model(
			    input_ids[:, :-1]	,      attention_mask=_SCREAMING_SNAKE_CASE	,      past_key_values=_SCREAMING_SNAKE_CASE	,      position_ids=_SCREAMING_SNAKE_CASE	,      )
			lowerCAmelCase__     :    Optional[Any]					  =      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,      dtype='''i4'''	)
			lowerCAmelCase__     :    List[Any]					  =      model(
			    input_ids[:, -1:]	,      past_key_values=outputs_cache.past_key_values	,      attention_mask=_SCREAMING_SNAKE_CASE	,      position_ids=_SCREAMING_SNAKE_CASE	,      )

			lowerCAmelCase__     :    Tuple					  =      model(_SCREAMING_SNAKE_CASE	,      attention_mask=_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Optional[int]					  =      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
			self.parent.assertTrue(diff < 1E-3	,      msg=F'Max diff is {diff}'	)





@require_flax
class       _a							(       _lowercase  ,		_lowercase  ,		unittest.TestCase):
	_a						:						Optional[int]   					=							(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
	_a						:						str   					=							(FlaxGPTaLMHeadModel,) if is_flax_available() else ()




	def 							UpperCAmelCase__(      self		:	Any	)-> Optional[Any]:
			lowerCAmelCase__     :    List[Any]					  =      FlaxGPTaModelTester(self	)




	def 							UpperCAmelCase__(      self		:	Tuple	)-> Optional[Any]:
			for model_class_name in self.all_model_classes:
					lowerCAmelCase__  ,			lowerCAmelCase__  ,			lowerCAmelCase__     :    Any					  =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.check_use_cache_forward(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Any	)-> str:
			for model_class_name in self.all_model_classes:
					lowerCAmelCase__  ,			lowerCAmelCase__  ,			lowerCAmelCase__     :    List[Any]					  =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.check_use_cache_forward_with_attn_mask(
					    _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)




	@slow
	def 							UpperCAmelCase__(      self		:	Dict	)-> List[str]:
			lowerCAmelCase__     :    str					  =      GPTaTokenizer.from_pretrained('''gpt2'''	,      pad_token='''</s>'''	,      padding_side='''left'''	)
			lowerCAmelCase__     :    Optional[Any]					  =      tokenizer(['''Hello this is a long string''', '''Hey''']	,      return_tensors='''np'''	,      padding=_SCREAMING_SNAKE_CASE	,      truncation=_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Tuple					  =      FlaxGPTaLMHeadModel.from_pretrained('''gpt2'''	)
			lowerCAmelCase__     :    Optional[Any]					  =      False
			lowerCAmelCase__     :    List[Any]					  =      model.config.eos_token_id

			lowerCAmelCase__     :    str					  =      jax.jit(model.generate	)

			lowerCAmelCase__     :    List[Any]					  =      jit_generate(inputs['''input_ids''']	,      attention_mask=inputs['''attention_mask''']	).sequences

			lowerCAmelCase__     :    Dict					  =      tokenizer.batch_decode(_SCREAMING_SNAKE_CASE	,      skip_special_tokens=_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Tuple					  =      [
			    '''Hello this is a long string of words. I\'m going to start with the first one.\n''',
			    '''Hey, I\'m not sure if I\'m going to be able to do''',
			]

			self.assertListEqual(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)




	@is_pt_flax_cross_test
	def 							UpperCAmelCase__(      self		:	List[str]	)-> int:
			lowerCAmelCase__  ,			lowerCAmelCase__     :    int					  =      self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__	):
							# prepare inputs
							lowerCAmelCase__     :    str					  =      self._prepare_for_class(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
							lowerCAmelCase__     :    Optional[int]					  =      {k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

							# load corresponding PyTorch class
							lowerCAmelCase__     :    List[str]					  =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
							lowerCAmelCase__     :    int					  =      getattr(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

							lowerCAmelCase__  ,			lowerCAmelCase__     :    Optional[Any]					  =      pt_inputs['''input_ids'''].shape
							lowerCAmelCase__     :    Tuple					  =      np.random.randint(0	,      seq_length - 1	,      size=(batch_size,)	)
							for batch_idx, start_index in enumerate(_SCREAMING_SNAKE_CASE	):
									lowerCAmelCase__     :    str					  =      0
									lowerCAmelCase__     :    Tuple					  =      1
									lowerCAmelCase__     :    int					  =      0
									lowerCAmelCase__     :    int					  =      1
							lowerCAmelCase__     :    Dict					  =      pt_model_class(_SCREAMING_SNAKE_CASE	).eval()
							lowerCAmelCase__     :    List[str]					  =      model_class(_SCREAMING_SNAKE_CASE	,      dtype=jnp.floataa	)

							lowerCAmelCase__     :    List[Any]					  =      convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,      _SCREAMING_SNAKE_CASE	)
							lowerCAmelCase__     :    Dict					  =      fx_state

							with torch.no_grad():
									lowerCAmelCase__     :    Optional[Any]					  =      pt_model(**_SCREAMING_SNAKE_CASE	).to_tuple()

							lowerCAmelCase__     :    Dict					  =      fx_model(**_SCREAMING_SNAKE_CASE	).to_tuple()
							self.assertEqual(len(_SCREAMING_SNAKE_CASE	)	,      len(_SCREAMING_SNAKE_CASE	)	,      '''Output lengths differ between Flax and PyTorch'''	)
							for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
									self.assert_almost_equals(fx_output[:, -1]	,      pt_output[:, -1].numpy()	,      4E-2	)

							with tempfile.TemporaryDirectory() as tmpdirname:
									pt_model.save_pretrained(_SCREAMING_SNAKE_CASE	)
									lowerCAmelCase__     :    Any					  =      model_class.from_pretrained(_SCREAMING_SNAKE_CASE	,      from_pt=_SCREAMING_SNAKE_CASE	)

							lowerCAmelCase__     :    int					  =      fx_model_loaded(**_SCREAMING_SNAKE_CASE	).to_tuple()
							self.assertEqual(
							    len(_SCREAMING_SNAKE_CASE	)	,      len(_SCREAMING_SNAKE_CASE	)	,      '''Output lengths differ between Flax and PyTorch'''	)
							for fx_output_loaded, pt_output in zip(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
									self.assert_almost_equals(fx_output_loaded[:, -1]	,      pt_output[:, -1].numpy()	,      4E-2	)




	@is_pt_flax_cross_test
	def 							UpperCAmelCase__(      self		:	List[str]	)-> Tuple:
			lowerCAmelCase__  ,			lowerCAmelCase__     :    Tuple					  =      self.model_tester.prepare_config_and_inputs_for_common()
			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__	):
							# prepare inputs
							lowerCAmelCase__     :    Dict					  =      self._prepare_for_class(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
							lowerCAmelCase__     :    str					  =      {k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

							# load corresponding PyTorch class
							lowerCAmelCase__     :    Any					  =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
							lowerCAmelCase__     :    Optional[int]					  =      getattr(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

							lowerCAmelCase__     :    str					  =      pt_model_class(_SCREAMING_SNAKE_CASE	).eval()
							lowerCAmelCase__     :    Tuple					  =      model_class(_SCREAMING_SNAKE_CASE	,      dtype=jnp.floataa	)

							lowerCAmelCase__     :    Dict					  =      load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE	,      fx_model.params	)
							lowerCAmelCase__  ,			lowerCAmelCase__     :    Dict					  =      pt_inputs['''input_ids'''].shape
							lowerCAmelCase__     :    List[str]					  =      np.random.randint(0	,      seq_length - 1	,      size=(batch_size,)	)
							for batch_idx, start_index in enumerate(_SCREAMING_SNAKE_CASE	):
									lowerCAmelCase__     :    Optional[int]					  =      0
									lowerCAmelCase__     :    Optional[Any]					  =      1
									lowerCAmelCase__     :    List[Any]					  =      0
									lowerCAmelCase__     :    Optional[Any]					  =      1

							# make sure weights are tied in PyTorch
							pt_model.tie_weights()

							with torch.no_grad():
									lowerCAmelCase__     :    int					  =      pt_model(**_SCREAMING_SNAKE_CASE	).to_tuple()

							lowerCAmelCase__     :    Dict					  =      fx_model(**_SCREAMING_SNAKE_CASE	).to_tuple()
							self.assertEqual(len(_SCREAMING_SNAKE_CASE	)	,      len(_SCREAMING_SNAKE_CASE	)	,      '''Output lengths differ between Flax and PyTorch'''	)
							for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
									self.assert_almost_equals(fx_output[:, -1]	,      pt_output[:, -1].numpy()	,      4E-2	)

							with tempfile.TemporaryDirectory() as tmpdirname:
									fx_model.save_pretrained(_SCREAMING_SNAKE_CASE	)
									lowerCAmelCase__     :    List[str]					  =      pt_model_class.from_pretrained(_SCREAMING_SNAKE_CASE	,      from_flax=_SCREAMING_SNAKE_CASE	)

							with torch.no_grad():
									lowerCAmelCase__     :    Any					  =      pt_model_loaded(**_SCREAMING_SNAKE_CASE	).to_tuple()

							self.assertEqual(
							    len(_SCREAMING_SNAKE_CASE	)	,      len(_SCREAMING_SNAKE_CASE	)	,      '''Output lengths differ between Flax and PyTorch'''	)
							for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
									self.assert_almost_equals(fx_output[:, -1]	,      pt_output[:, -1].numpy()	,      4E-2	)




	@slow
	def 							UpperCAmelCase__(      self		:	Union[str, Any]	)-> Tuple:
			for model_class_name in self.all_model_classes:
					lowerCAmelCase__     :    Union[str, Any]					  =      model_class_name.from_pretrained('''gpt2'''	,      from_pt=_SCREAMING_SNAKE_CASE	)
					lowerCAmelCase__     :    int					  =      model(np.ones((1, 1)	)	)
					self.assertIsNotNone(_SCREAMING_SNAKE_CASE	)








import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase		=	logging.get_logger(__name__)

lowerCamelCase		=	{
    '''facebook/sam-vit-huge''': '''https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json''',
    '''facebook/sam-vit-large''': '''https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json''',
    '''facebook/sam-vit-base''': '''https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json''',
}





class       _a							(       _lowercase):




	def __init__(      self		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple=256	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=1024	,      _SCREAMING_SNAKE_CASE		:	List[Any]=16	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=16	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=4	,      _SCREAMING_SNAKE_CASE		:	Any="gelu"	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=1E-6	,      **_SCREAMING_SNAKE_CASE		:	List[Any]	,      )-> Tuple:
			super().__init__(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      hidden_size
			lowerCAmelCase__     :    Tuple					  =      image_size
			lowerCAmelCase__     :    List[Any]					  =      patch_size
			lowerCAmelCase__     :    Any					  =      image_size // patch_size
			lowerCAmelCase__     :    Dict					  =      mask_input_channels
			lowerCAmelCase__     :    Dict					  =      num_point_embeddings
			lowerCAmelCase__     :    str					  =      hidden_act
			lowerCAmelCase__     :    List[Any]					  =      layer_norm_eps





class       _a							(       _lowercase):




	def __init__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=256	,      _SCREAMING_SNAKE_CASE		:	Dict="relu"	,      _SCREAMING_SNAKE_CASE		:	str=2048	,      _SCREAMING_SNAKE_CASE		:	List[Any]=2	,      _SCREAMING_SNAKE_CASE		:	Dict=8	,      _SCREAMING_SNAKE_CASE		:	str=2	,      _SCREAMING_SNAKE_CASE		:	Tuple=3	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=3	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=256	,      _SCREAMING_SNAKE_CASE		:	Dict=1E-6	,      **_SCREAMING_SNAKE_CASE		:	Any	,      )-> List[Any]:
			super().__init__(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[str]					  =      hidden_size
			lowerCAmelCase__     :    Union[str, Any]					  =      hidden_act
			lowerCAmelCase__     :    str					  =      mlp_dim
			lowerCAmelCase__     :    Optional[int]					  =      num_hidden_layers
			lowerCAmelCase__     :    int					  =      num_attention_heads
			lowerCAmelCase__     :    List[str]					  =      attention_downsample_rate
			lowerCAmelCase__     :    str					  =      num_multimask_outputs
			lowerCAmelCase__     :    Optional[int]					  =      iou_head_depth
			lowerCAmelCase__     :    Union[str, Any]					  =      iou_head_hidden_dim
			lowerCAmelCase__     :    Tuple					  =      layer_norm_eps





class       _a							(       _lowercase):




	def __init__(      self		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	int=768	,      _SCREAMING_SNAKE_CASE		:	str=256	,      _SCREAMING_SNAKE_CASE		:	int=12	,      _SCREAMING_SNAKE_CASE		:	Tuple=12	,      _SCREAMING_SNAKE_CASE		:	Dict=3	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=1024	,      _SCREAMING_SNAKE_CASE		:	Tuple=16	,      _SCREAMING_SNAKE_CASE		:	List[Any]="gelu"	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=1E-06	,      _SCREAMING_SNAKE_CASE		:	int=0.0	,      _SCREAMING_SNAKE_CASE		:	int=1E-10	,      _SCREAMING_SNAKE_CASE		:	List[Any]=True	,      _SCREAMING_SNAKE_CASE		:	Any=4.0	,      _SCREAMING_SNAKE_CASE		:	List[str]=True	,      _SCREAMING_SNAKE_CASE		:	int=True	,      _SCREAMING_SNAKE_CASE		:	List[str]=14	,      _SCREAMING_SNAKE_CASE		:	Dict=[2, 5, 8, 11]	,      _SCREAMING_SNAKE_CASE		:	List[Any]=128	,      _SCREAMING_SNAKE_CASE		:	Dict=None	,      **_SCREAMING_SNAKE_CASE		:	str	,      )-> Any:
			super().__init__(**_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Optional[Any]					  =      hidden_size
			lowerCAmelCase__     :    Any					  =      output_channels
			lowerCAmelCase__     :    Optional[int]					  =      num_hidden_layers
			lowerCAmelCase__     :    Dict					  =      num_attention_heads
			lowerCAmelCase__     :    List[Any]					  =      num_channels
			lowerCAmelCase__     :    int					  =      image_size
			lowerCAmelCase__     :    Tuple					  =      patch_size
			lowerCAmelCase__     :    Union[str, Any]					  =      hidden_act
			lowerCAmelCase__     :    List[str]					  =      layer_norm_eps
			lowerCAmelCase__     :    Any					  =      attention_dropout
			lowerCAmelCase__     :    int					  =      initializer_range
			lowerCAmelCase__     :    List[Any]					  =      qkv_bias
			lowerCAmelCase__     :    List[Any]					  =      mlp_ratio
			lowerCAmelCase__     :    List[Any]					  =      use_abs_pos
			lowerCAmelCase__     :    Union[str, Any]					  =      use_rel_pos
			lowerCAmelCase__     :    Optional[Any]					  =      window_size
			lowerCAmelCase__     :    Union[str, Any]					  =      global_attn_indexes
			lowerCAmelCase__     :    Union[str, Any]					  =      num_pos_feats
			lowerCAmelCase__     :    str					  =      int(hidden_size * mlp_ratio	) if mlp_dim is None else mlp_dim





class       _a							(       _lowercase):
	_a						:						Tuple   					=							'''sam'''
	_a						:						Optional[Any]   					=							True




	def __init__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Any=None	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=None	,      _SCREAMING_SNAKE_CASE		:	int=None	,      _SCREAMING_SNAKE_CASE		:	Tuple=0.02	,      **_SCREAMING_SNAKE_CASE		:	int	,      )-> Optional[int]:
			super().__init__(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      vision_config if vision_config is not None else {}
			lowerCAmelCase__     :    List[str]					  =      prompt_encoder_config if prompt_encoder_config is not None else {}
			lowerCAmelCase__     :    List[str]					  =      mask_decoder_config if mask_decoder_config is not None else {}

			if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
					lowerCAmelCase__     :    Any					  =      vision_config.to_dict()
			if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
					lowerCAmelCase__     :    Any					  =      prompt_encoder_config.to_dict()
			if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
					lowerCAmelCase__     :    List[Any]					  =      mask_decoder_config.to_dict()

			lowerCAmelCase__     :    Any					  =      SamVisionConfig(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[Any]					  =      SamPromptEncoderConfig(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Dict					  =      SamMaskDecoderConfig(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[Any]					  =      initializer_range




	def 							UpperCAmelCase__(      self		:	Any	)-> int:
			lowerCAmelCase__     :    Optional[int]					  =      copy.deepcopy(self.__dict__	)
			lowerCAmelCase__     :    Any					  =      self.vision_config.to_dict()
			lowerCAmelCase__     :    Optional[Any]					  =      self.prompt_encoder_config.to_dict()
			lowerCAmelCase__     :    str					  =      self.mask_decoder_config.to_dict()
			lowerCAmelCase__     :    Any					  =      self.__class__.model_type
			return output




