



import copy
import os
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union


if TYPE_CHECKING:
	from ...processing_utils import ProcessorMixin
	from ...utils import TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A							       =		logging.get_logger(__name__)

__A							       =		{
    '''google/owlvit-base-patch32''': '''https://huggingface.co/google/owlvit-base-patch32/resolve/main/config.json''',
    '''google/owlvit-base-patch16''': '''https://huggingface.co/google/owlvit-base-patch16/resolve/main/config.json''',
    '''google/owlvit-large-patch14''': '''https://huggingface.co/google/owlvit-large-patch14/resolve/main/config.json''',
}






class 						lowercase		(					snake_case__):






	"""simple docstring"""


	a__ :   str						  =				'owlvit_text_model'




	def __init__(      self      :		int		,    __UpperCAmelCase      :		Any=49_408		,    __UpperCAmelCase      :		Dict=512		,    __UpperCAmelCase      :		Optional[Any]=2_048		,    __UpperCAmelCase      :		Any=12		,    __UpperCAmelCase      :		Union[str, Any]=8		,    __UpperCAmelCase      :		Optional[int]=16		,    __UpperCAmelCase      :		Optional[Any]="quick_gelu"		,    __UpperCAmelCase      :		Optional[Any]=1E-5		,    __UpperCAmelCase      :		List[str]=0.0		,    __UpperCAmelCase      :		Union[str, Any]=0.02		,    __UpperCAmelCase      :		Tuple=1.0		,    __UpperCAmelCase      :		List[str]=0		,    __UpperCAmelCase      :		List[Any]=49_406		,    __UpperCAmelCase      :		Optional[Any]=49_407		,    **__UpperCAmelCase      :		List[Any]		,    )     ->					Union[str, Any]:
							super().__init__(pad_token_id=__UpperCAmelCase		,    bos_token_id=__UpperCAmelCase		,    eos_token_id=__UpperCAmelCase		,    **__UpperCAmelCase	)

							UpperCAmelCase_=  vocab_size
							UpperCAmelCase_=  hidden_size
							UpperCAmelCase_=  intermediate_size
							UpperCAmelCase_=  num_hidden_layers
							UpperCAmelCase_=  num_attention_heads
							UpperCAmelCase_=  max_position_embeddings
							UpperCAmelCase_=  hidden_act
							UpperCAmelCase_=  layer_norm_eps
							UpperCAmelCase_=  attention_dropout
							UpperCAmelCase_=  initializer_range
							UpperCAmelCase_=  initializer_factor




	@classmethod
	def     _SCREAMING_SNAKE_CASE						(      cls      :		Optional[int]		,    __UpperCAmelCase      :		List[str]		,    **__UpperCAmelCase      :		Optional[int]	)     ->					"PretrainedConfig":
							cls._set_token_in_kwargs(__UpperCAmelCase	)

							UpperCAmelCase_,  UpperCAmelCase_=  cls.get_config_dict(__UpperCAmelCase		,    **__UpperCAmelCase	)

							# get the text config dict if we are loading from OwlViTConfig
							if config_dict.get("""model_type"""	) == "owlvit":
													UpperCAmelCase_=  config_dict["""text_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		(					snake_case__):






	"""simple docstring"""


	a__ :   Dict						  =				'owlvit_vision_model'




	def __init__(      self      :		int		,    __UpperCAmelCase      :		Optional[Any]=768		,    __UpperCAmelCase      :		Any=3_072		,    __UpperCAmelCase      :		Dict=12		,    __UpperCAmelCase      :		List[str]=12		,    __UpperCAmelCase      :		Tuple=3		,    __UpperCAmelCase      :		Tuple=768		,    __UpperCAmelCase      :		Dict=32		,    __UpperCAmelCase      :		Union[str, Any]="quick_gelu"		,    __UpperCAmelCase      :		Tuple=1E-5		,    __UpperCAmelCase      :		Optional[Any]=0.0		,    __UpperCAmelCase      :		Optional[Any]=0.02		,    __UpperCAmelCase      :		Optional[Any]=1.0		,    **__UpperCAmelCase      :		Union[str, Any]		,    )     ->					Union[str, Any]:
							super().__init__(**__UpperCAmelCase	)

							UpperCAmelCase_=  hidden_size
							UpperCAmelCase_=  intermediate_size
							UpperCAmelCase_=  num_hidden_layers
							UpperCAmelCase_=  num_attention_heads
							UpperCAmelCase_=  num_channels
							UpperCAmelCase_=  image_size
							UpperCAmelCase_=  patch_size
							UpperCAmelCase_=  hidden_act
							UpperCAmelCase_=  layer_norm_eps
							UpperCAmelCase_=  attention_dropout
							UpperCAmelCase_=  initializer_range
							UpperCAmelCase_=  initializer_factor




	@classmethod
	def     _SCREAMING_SNAKE_CASE						(      cls      :		List[str]		,    __UpperCAmelCase      :		Tuple		,    **__UpperCAmelCase      :		Optional[int]	)     ->					"PretrainedConfig":
							cls._set_token_in_kwargs(__UpperCAmelCase	)

							UpperCAmelCase_,  UpperCAmelCase_=  cls.get_config_dict(__UpperCAmelCase		,    **__UpperCAmelCase	)

							# get the vision config dict if we are loading from OwlViTConfig
							if config_dict.get("""model_type"""	) == "owlvit":
													UpperCAmelCase_=  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		(					snake_case__):






	"""simple docstring"""


	a__ :   Optional[int]						  =				'owlvit'
	a__ :   List[str]						  =				True




	def __init__(      self      :		Tuple		,    __UpperCAmelCase      :		Optional[Any]=None		,    __UpperCAmelCase      :		str=None		,    __UpperCAmelCase      :		List[Any]=512		,    __UpperCAmelCase      :		List[str]=2.6_592		,    __UpperCAmelCase      :		Dict=True		,    **__UpperCAmelCase      :		List[Any]		,    )     ->					Optional[int]:
							super().__init__(**__UpperCAmelCase	)

							if text_config is None:
													UpperCAmelCase_=  {}
													logger.info("""text_config is None. Initializing the OwlViTTextConfig with default values."""	)

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

							UpperCAmelCase_=  OwlViTTextConfig(**__UpperCAmelCase	)
							UpperCAmelCase_=  OwlViTVisionConfig(**__UpperCAmelCase	)

							UpperCAmelCase_=  projection_dim
							UpperCAmelCase_=  logit_scale_init_value
							UpperCAmelCase_=  return_dict
							UpperCAmelCase_=  1.0




	@classmethod
	def     _SCREAMING_SNAKE_CASE						(      cls      :		List[str]		,    __UpperCAmelCase      :		List[str]		,    **__UpperCAmelCase      :		Optional[Any]	)     ->					"PretrainedConfig":
							cls._set_token_in_kwargs(__UpperCAmelCase	)

							UpperCAmelCase_,  UpperCAmelCase_=  cls.get_config_dict(__UpperCAmelCase		,    **__UpperCAmelCase	)

							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	)




	@classmethod
	def     _SCREAMING_SNAKE_CASE						(      cls      :		str		,    __UpperCAmelCase      :		List[str]		,    __UpperCAmelCase      :		Optional[int]		,    **__UpperCAmelCase      :		Union[str, Any]	)     ->					Any:
							UpperCAmelCase_=  {}
							UpperCAmelCase_=  text_config
							UpperCAmelCase_=  vision_config

							return cls.from_dict(__UpperCAmelCase		,    **__UpperCAmelCase	)




	def     _SCREAMING_SNAKE_CASE						(      self      :		List[str]	)     ->					Tuple:
							UpperCAmelCase_=  copy.deepcopy(self.__dict__	)
							UpperCAmelCase_=  self.text_config.to_dict()
							UpperCAmelCase_=  self.vision_config.to_dict()
							UpperCAmelCase_=  self.__class__.model_type
							return output






class 						lowercase		(					snake_case__):






	"""simple docstring"""


	@property
	def     _SCREAMING_SNAKE_CASE						(      self      :		Optional[int]	)     ->					Mapping[str, Mapping[int, str]]:
							return OrderedDict(
							    [
							        ("""input_ids""", {0: """batch""", 1: """sequence"""}),
							        ("""pixel_values""", {0: """batch""", 1: """num_channels""", 2: """height""", 3: """width"""}),
							        ("""attention_mask""", {0: """batch""", 1: """sequence"""}),
							    ]	)




	@property
	def     _SCREAMING_SNAKE_CASE						(      self      :		str	)     ->					Mapping[str, Mapping[int, str]]:
							return OrderedDict(
							    [
							        ("""logits_per_image""", {0: """batch"""}),
							        ("""logits_per_text""", {0: """batch"""}),
							        ("""text_embeds""", {0: """batch"""}),
							        ("""image_embeds""", {0: """batch"""}),
							    ]	)




	@property
	def     _SCREAMING_SNAKE_CASE						(      self      :		int	)     ->					float:
							return 1E-4




	def     _SCREAMING_SNAKE_CASE						(      self      :		Optional[Any]		,    __UpperCAmelCase      :		List[str]		,    __UpperCAmelCase      :		List[Any] = -1		,    __UpperCAmelCase      :		Dict = -1		,    __UpperCAmelCase      :		int = None		,    )     ->					Mapping[str, Any]:
							UpperCAmelCase_=  super().generate_dummy_inputs(
							    processor.tokenizer		,    batch_size=__UpperCAmelCase		,    seq_length=__UpperCAmelCase		,    framework=__UpperCAmelCase	)
							UpperCAmelCase_=  super().generate_dummy_inputs(
							    processor.image_processor		,    batch_size=__UpperCAmelCase		,    framework=__UpperCAmelCase	)
							return {**text_input_dict, **image_input_dict}




	@property
	def     _SCREAMING_SNAKE_CASE						(      self      :		Dict	)     ->					int:
							return 14









'''simple docstring'''





import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import DecisionTransformerModel
  from transformers.models.decision_transformer.modeling_decision_transformer import (
      DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )






class __A     :



 '''simple docstring'''







 def __init__(self					,  A					,  A=13					,  A=7					,  A=6					,  A=17					,  A=23					,  A=11					,  A=True					,  )       ->	Tuple:


      """simple docstring"""



      _a           =    parent
      _a           =    batch_size
      _a           =    seq_length
      _a           =    act_dim
      _a           =    state_dim
      _a           =    hidden_size
      _a           =    max_length
      _a           =    is_training



 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      _a           =    floats_tensor((self.batch_size, self.seq_length, self.state_dim)			)
      _a           =    floats_tensor((self.batch_size, self.seq_length, self.act_dim)			)
      _a           =    floats_tensor((self.batch_size, self.seq_length, 1)			)
      _a           =    floats_tensor((self.batch_size, self.seq_length, 1)			)
      _a           =    ids_tensor((self.batch_size, self.seq_length)					,  vocab_size=1_000			)
      _a           =    random_attention_mask((self.batch_size, self.seq_length)			)

      _a           =    self.get_config()

      return (
          config,
          states,
          actions,
          rewards,
          returns_to_go,
          timesteps,
          attention_mask,
      )



 def    a__    (self			)       ->	str:


      """simple docstring"""



      return DecisionTransformerConfig(
          batch_size=self.batch_size					,  seq_length=self.seq_length					,  act_dim=self.act_dim					,  state_dim=self.state_dim					,  hidden_size=self.hidden_size					,  max_length=self.max_length					,  )



 def    a__    (self					,  A					,  A					,  A					,  A					,  A					,  A					,  A					,  )       ->	List[Any]:


      """simple docstring"""



      _a           =    DecisionTransformerModel(config=A			)
      model.to(A			)
      model.eval()
      _a           =    model(A					,  A					,  A					,  A					,  A					,  A			)

      self.parent.assertEqual(result.state_preds.shape					,  states.shape			)
      self.parent.assertEqual(result.action_preds.shape					,  actions.shape			)
      self.parent.assertEqual(result.return_preds.shape					,  returns_to_go.shape			)
      self.parent.assertEqual(
          result.last_hidden_state.shape					,  (self.batch_size, self.seq_length * 3, self.hidden_size)			)  # seq length *3 as there are 3 modelities: states, returns and actions



 def    a__    (self			)       ->	Dict:


      """simple docstring"""



      _a           =    self.prepare_config_and_inputs()
      (
          (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       
      )           =    config_and_inputs
      _a           =    {
          '''states''': states,
          '''actions''': actions,
          '''rewards''': rewards,
          '''returns_to_go''': returns_to_go,
          '''timesteps''': timesteps,
          '''attention_mask''': attention_mask,
      }
      return config, inputs_dict






@require_torch
class __A     (    A  ,  A  ,  A  ,  unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase       :			Optional[Any]		=			(DecisionTransformerModel,) if is_torch_available() else ()
 __lowerCamelCase       :			List[str]		=			()
 __lowerCamelCase       :			Tuple		=			{'feature-extraction': DecisionTransformerModel} if is_torch_available() else {}

 # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
 __lowerCamelCase       :			str		=			False

 # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
 __lowerCamelCase       :			List[str]		=			False
 __lowerCamelCase       :			List[str]		=			False
 __lowerCamelCase       :			Tuple		=			False
 __lowerCamelCase       :			str		=			False
 __lowerCamelCase       :			Dict		=			False
 __lowerCamelCase       :			Tuple		=			False
 __lowerCamelCase       :			Tuple		=			False
 __lowerCamelCase       :			Dict		=			False
 __lowerCamelCase       :			List[str]		=			False



 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      _a           =    DecisionTransformerModelTester(self			)
      _a           =    ConfigTester(self					,  config_class=A					,  hidden_size=37			)



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


      """simple docstring"""



      self.config_tester.run_common_tests()



 def    a__    (self			)       ->	List[Any]:


      """simple docstring"""



      _a           =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A			)



 @slow
 def    a__    (self			)       ->	Optional[Any]:


      """simple docstring"""



      for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _a           =    DecisionTransformerModel.from_pretrained(A			)
           self.assertIsNotNone(A			)



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


      """simple docstring"""



      _a    ,       _a           =    self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _a           =    model_class(A			)
           _a           =    inspect.signature(model.forward			)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           _a           =    [*signature.parameters.keys()]

           _a           =    [
               '''states''',
               '''actions''',
               '''rewards''',
               '''returns_to_go''',
               '''timesteps''',
               '''attention_mask''',
           ]

           self.assertListEqual(arg_names[: len(A			)]					,  A			)






@require_torch
class __A     (    unittest.TestCase		):



 '''simple docstring'''







 @slow
 def    a__    (self			)       ->	Optional[Any]:


      """simple docstring"""



      _a           =    2  # number of steps of autoregressive prediction we will perform
      _a           =    10  # defined by the RL environment, may be normalized
      _a           =    DecisionTransformerModel.from_pretrained('''edbeeching/decision-transformer-gym-hopper-expert'''			)
      _a           =    model.to(A			)
      _a           =    model.config
      torch.manual_seed(0			)
      _a           =    torch.randn(1					,  1					,  config.state_dim			).to(device=A					,  dtype=torch.floataa			)  # env.reset()

      _a           =    torch.tensor(
          [[0.242793, -0.28693074, 0.8742613], [0.67815274, -0.08101085, -0.12952147]]					,  device=A			)

      _a           =    torch.tensor(A					,  device=A					,  dtype=torch.floataa			).reshape(1					,  1					,  1			)
      _a           =    state
      _a           =    torch.zeros(1					,  0					,  config.act_dim					,  device=A					,  dtype=torch.floataa			)
      _a           =    torch.zeros(1					,  0					,  device=A					,  dtype=torch.floataa			)
      _a           =    torch.tensor(0					,  device=A					,  dtype=torch.long			).reshape(1					,  1			)

      for step in range(A			):
           _a           =    torch.cat([actions, torch.zeros(1					,  1					,  config.act_dim					,  device=A			)]					,  dim=1			)
           _a           =    torch.cat([rewards, torch.zeros(1					,  1					,  device=A			)]					,  dim=1			)

           _a           =    torch.ones(1					,  states.shape[1]			).to(dtype=torch.long					,  device=states.device			)

           with torch.no_grad():
                _a    ,       _a    ,       _a           =    model(
                    states=A					,  actions=A					,  rewards=A					,  returns_to_go=A					,  timesteps=A					,  attention_mask=A					,  return_dict=A					,  )

           self.assertEqual(action_pred.shape					,  actions.shape			)
           self.assertTrue(torch.allclose(action_pred[0, -1]					,  expected_outputs[step]					,  atol=1E-4			)			)
           _a    ,       _a    ,       _a    ,       _a           =    (  # env.step(action)
               torch.randn(1					,  1					,  config.state_dim			).to(device=A					,  dtype=torch.floataa			),
               1.0,
               False,
               {},
           )

           _a           =    action_pred[0, -1]
           _a           =    torch.cat([states, state]					,  dim=1			)
           _a           =    returns_to_go[0, -1] - reward
           _a           =    torch.cat([returns_to_go, pred_return.reshape(1					,  1					,  1			)]					,  dim=1			)
           _a           =    torch.cat(
               [timesteps, torch.ones((1, 1)					,  device=A					,  dtype=torch.long			) * (step + 1)]					,  dim=1			)



