




import inspect
import unittest

from transformers import ViTHybridConfig
from transformers.testing_utils import require_accelerate, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import ViTHybridForImageClassification, ViTHybridImageProcessor, ViTHybridModel
  from transformers.models.vit_hybrid.modeling_vit_hybrid import VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image
class 			_UpperCamelCase      :
      def __init__(			self   :List[str] ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :Tuple=13 ,      lowerCamelCase   :Any=64 ,      lowerCamelCase   :int=2 ,      lowerCamelCase   :Dict=3 ,      lowerCamelCase   :List[str]=True ,      lowerCamelCase   :Union[str, Any]=True ,      lowerCamelCase   :Optional[int]=32 ,      lowerCamelCase   :Optional[int]=5 ,      lowerCamelCase   :int=4 ,      lowerCamelCase   :Any=37 ,      lowerCamelCase   :int="gelu" ,      lowerCamelCase   :int=0.1 ,      lowerCamelCase   :Any=0.1 ,      lowerCamelCase   :str=10 ,      lowerCamelCase   :Any=0.02 ,      lowerCamelCase   :str=[1, 16, 4, 4] ,      lowerCamelCase   :List[str]=None ,      )     ->			int:
       UpperCAmelCase__							 =						parent
       UpperCAmelCase__							 =						batch_size
       UpperCAmelCase__							 =						image_size
       UpperCAmelCase__							 =						patch_size
       UpperCAmelCase__							 =						num_channels
       UpperCAmelCase__							 =						is_training
       UpperCAmelCase__							 =						use_labels
       UpperCAmelCase__							 =						hidden_size
       UpperCAmelCase__							 =						num_hidden_layers
       UpperCAmelCase__							 =						num_attention_heads
       UpperCAmelCase__							 =						intermediate_size
       UpperCAmelCase__							 =						hidden_act
       UpperCAmelCase__							 =						hidden_dropout_prob
       UpperCAmelCase__							 =						attention_probs_dropout_prob
       UpperCAmelCase__							 =						type_sequence_label_size
       UpperCAmelCase__							 =						initializer_range
       UpperCAmelCase__							 =						scope
       UpperCAmelCase__							 =						backbone_featmap_shape

       # in ViT hybrid, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       # the number of patches is based on the feature map of the backbone, which by default uses an output stride
       # of 32, which means that the feature map has a spatial resolution of 1/32 of the input image size
       UpperCAmelCase__							 =						(self.image_size // 32) ** 2
       UpperCAmelCase__							 =						num_patches + 1
      def 	UpperCAmelCase_						(			self   :List[str]      )     ->			Tuple:
       UpperCAmelCase__							 =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

       UpperCAmelCase__							 =						None
       if self.use_labels:
        UpperCAmelCase__							 =						ids_tensor([self.batch_size] ,      self.type_sequence_label_size      )

       UpperCAmelCase__							 =						self.get_config()

       return config, pixel_values, labels
      def 	UpperCAmelCase_						(			self   :Dict      )     ->			int:
       UpperCAmelCase__							 =						{
           "global_padding": "same",
           "layer_type": "bottleneck",
           "depths": [3, 4, 9],
           "out_features": ["stage1", "stage2", "stage3"],
           "embedding_dynamic_padding": True,
           "hidden_sizes": [4, 8, 16, 32],
           "num_groups": 2,
       }

       return ViTHybridConfig(
           image_size=self.image_size ,      patch_size=self.patch_size ,      num_channels=self.num_channels ,      hidden_size=self.hidden_size ,      num_hidden_layers=self.num_hidden_layers ,      num_attention_heads=self.num_attention_heads ,      intermediate_size=self.intermediate_size ,      hidden_act=self.hidden_act ,      hidden_dropout_prob=self.hidden_dropout_prob ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob ,      is_decoder=lowerCamelCase ,      initializer_range=self.initializer_range ,      backbone_featmap_shape=self.backbone_featmap_shape ,      backbone_config=lowerCamelCase ,      )
      def 	UpperCAmelCase_						(			self   :List[str] ,      lowerCamelCase   :Optional[int] ,      lowerCamelCase   :str ,      lowerCamelCase   :List[str]      )     ->			Dict:
       UpperCAmelCase__							 =						ViTHybridModel(config=lowerCamelCase      )
       model.to(lowerCamelCase      )
       model.eval()
       UpperCAmelCase__							 =						model(lowerCamelCase      )
       self.parent.assertEqual(result.last_hidden_state.shape ,      (self.batch_size, self.seq_length, self.hidden_size)      )
      def 	UpperCAmelCase_						(			self   :Any ,      lowerCamelCase   :int ,      lowerCamelCase   :List[str] ,      lowerCamelCase   :List[str]      )     ->			Optional[Any]:
       UpperCAmelCase__							 =						self.type_sequence_label_size
       UpperCAmelCase__							 =						ViTHybridForImageClassification(lowerCamelCase      )
       model.to(lowerCamelCase      )
       model.eval()
       UpperCAmelCase__							 =						model(lowerCamelCase ,      labels=lowerCamelCase      )
       self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.type_sequence_label_size)      )







      def 	UpperCAmelCase_						(			self   :Tuple      )     ->			Any:
       UpperCAmelCase__							 =						self.prepare_config_and_inputs()
       UpperCAmelCase__ , UpperCAmelCase__ , UpperCAmelCase__							 =						config_and_inputs
       UpperCAmelCase__							 =						{"pixel_values": pixel_values}
       return config, inputs_dict



@require_torch
class 			_UpperCamelCase      (				lowerCAmelCase	,       lowerCAmelCase	,       unittest.TestCase		):
      UpperCAmelCase_        =		(ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
      UpperCAmelCase_        =		(
          {"""feature-extraction""": ViTHybridModel, """image-classification""": ViTHybridForImageClassification}
          if is_torch_available()
          else {}
      )
      UpperCAmelCase_        =		False
      UpperCAmelCase_        =		False
      UpperCAmelCase_        =		False
      def 	UpperCAmelCase_						(			self   :Union[str, Any]      )     ->			Union[str, Any]:
       UpperCAmelCase__							 =						ViTHybridModelTester(self      )
       UpperCAmelCase__							 =						ConfigTester(self ,      config_class=lowerCamelCase ,      has_text_modality=lowerCamelCase ,      hidden_size=37      )
      def 	UpperCAmelCase_						(			self   :Optional[Any]      )     ->			int:
       self.config_tester.run_common_tests()
      @unittest.skip(reason="ViT does not use inputs_embeds"      )
      def 	UpperCAmelCase_						(			self   :str      )     ->			Optional[Any]:
       pass
      def 	UpperCAmelCase_						(			self   :Any      )     ->			int:
       UpperCAmelCase__ , UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCAmelCase__							 =						model_class(lowerCamelCase      )
        self.assertIsInstance(model.get_input_embeddings() ,      (nn.Module)      )
        UpperCAmelCase__							 =						model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(lowerCamelCase ,      nn.Linear      )      )
      def 	UpperCAmelCase_						(			self   :Optional[int]      )     ->			str:
       UpperCAmelCase__ , UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs_for_common()

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

        UpperCAmelCase__							 =						["pixel_values"]
        self.assertListEqual(arg_names[:1] ,      lowerCamelCase      )
      def 	UpperCAmelCase_						(			self   :Union[str, Any]      )     ->			Dict:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*lowerCamelCase      )
      def 	UpperCAmelCase_						(			self   :Optional[int]      )     ->			Optional[Any]:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*lowerCamelCase      )
      def 	UpperCAmelCase_						(			self   :List[Any]      )     ->			Any:
       UpperCAmelCase__ , UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs_for_common()

       UpperCAmelCase__							 =						_config_zero_init(lowerCamelCase      )
       for model_class in self.all_model_classes:
        UpperCAmelCase__							 =						model_class(config=lowerCamelCase      )
        # Skip the check for the backbone
        for name, module in model.named_modules():
         if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
          UpperCAmelCase__							 =						[f'''{name}.{key}''' for key in module.state_dict().keys()]
          break

        for name, param in model.named_parameters():
         if param.requires_grad:
          if name in backbone_params:
           continue
          self.assertIn(
              ((param.data.mean() * 1e9).round() / 1e9).item() ,      [0.0, 1.0] ,      msg=f'''Parameter {name} of model {model_class} seems not properly initialized''' ,      )







      @slow
      def 	UpperCAmelCase_						(			self   :Any      )     ->			str:
       for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							 =						ViTHybridModel.from_pretrained(lowerCamelCase      )
        self.assertIsNotNone(lowerCamelCase      )



def 						lowerCAmelCase		(		):




 """simple docstring"""







 UpperCAmelCase__							 =						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
 return image


@require_torch
@require_vision
class 			_UpperCamelCase      (				unittest.TestCase		):
      @cached_property
      def 	UpperCAmelCase_						(			self   :Tuple      )     ->			Optional[Any]:
       return (
           ViTHybridImageProcessor.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]      )
           if is_vision_available()
           else None
       )
      @slow
      def 	UpperCAmelCase_						(			self   :List[Any]      )     ->			int:
       UpperCAmelCase__							 =						ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]      ).to(
           lowerCamelCase      )

       UpperCAmelCase__							 =						self.default_image_processor
       UpperCAmelCase__							 =						prepare_img()
       UpperCAmelCase__							 =						image_processor(images=lowerCamelCase ,      return_tensors="pt"      ).to(lowerCamelCase      )

       # forward pass
       with torch.no_grad():
        UpperCAmelCase__							 =						model(**lowerCamelCase      )

       # verify the logits
       UpperCAmelCase__							 =						torch.Size((1, 1000)      )
       self.assertEqual(outputs.logits.shape ,      lowerCamelCase      )

       UpperCAmelCase__							 =						torch.tensor([-1.90_90, -0.49_93, -0.23_89]      ).to(lowerCamelCase      )

       self.assertTrue(torch.allclose(outputs.logits[0, :3] ,      lowerCamelCase ,      atol=1e-4      )      )







      @slow
      @require_accelerate
      def 	UpperCAmelCase_						(			self   :List[str]      )     ->			List[str]:
       UpperCAmelCase__							 =						ViTHybridImageProcessor.from_pretrained("google/vit-hybrid-base-bit-384"      )
       UpperCAmelCase__							 =						ViTHybridForImageClassification.from_pretrained("google/vit-hybrid-base-bit-384" ,      device_map="auto"      )

       UpperCAmelCase__							 =						prepare_img()

       UpperCAmelCase__							 =						image_processor(images=lowerCamelCase ,      return_tensors="pt"      )
       UpperCAmelCase__							 =						model(**lowerCamelCase      )
       UpperCAmelCase__							 =						outputs.logits
       # model predicts one of the 1000 ImageNet classes
       UpperCAmelCase__							 =						logits.argmax(-1      ).item()

       self.assertTrue(model.config.idalabel[predicted_class_idx] ,      "tabby, tabby cat"      )











import torch

from transformers import AutoModel
class 			_UpperCamelCase      (				torch.nn.Module		):
      def __init__(			self   :str ,      lowerCamelCase   :Tuple="sayef/fsner-bert-base-uncased"      )     ->			int:
       super(lowerCamelCase ,      self      ).__init__()

       UpperCAmelCase__							 =						AutoModel.from_pretrained(lowerCamelCase ,      return_dict=lowerCamelCase      )
       UpperCAmelCase__							 =						torch.nn.CosineSimilarity(3 ,      1e-08      )
       UpperCAmelCase__							 =						torch.nn.Softmax(dim=1      )
      def 	UpperCAmelCase_						(			self   :Union[str, Any] ,      **lowerCamelCase   :Tuple      )     ->			Dict:
       return self.bert(**lowerCamelCase      ).last_hidden_state
      def 	UpperCAmelCase_						(			self   :Any ,      lowerCamelCase   :Union[str, Any]      )     ->			Union[str, Any]:
       return token_embeddings.sum(2 ,      keepdim=lowerCamelCase      )
      def 	UpperCAmelCase_						(			self   :Dict ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :int ,      lowerCamelCase   :Union[str, Any]=1      )     ->			Dict:
       return self.softmax(T * self.cos(lowerCamelCase ,      lowerCamelCase      )      )







      def 	UpperCAmelCase_						(			self   :Union[str, Any] ,      lowerCamelCase   :Any ,      lowerCamelCase   :Any      )     ->			Union[str, Any]:
       UpperCAmelCase__							 =						W_supports["sizes"].tolist()
       UpperCAmelCase__							 =						W_supports["start_token_id"].item()
       UpperCAmelCase__							 =						W_supports["end_token_id"].item()

       del W_supports["sizes"]
       del W_supports["start_token_id"]
       del W_supports["end_token_id"]

       UpperCAmelCase__							 =						self.BERT(**lowerCamelCase      )
       UpperCAmelCase__							 =						self.BERT(**lowerCamelCase      )

       UpperCAmelCase__							 =						None
       UpperCAmelCase__							 =						None

       UpperCAmelCase__							 =						W_supports["input_ids"] == start_token_id
       UpperCAmelCase__							 =						W_supports["input_ids"] == end_token_id

       for i, size in enumerate(lowerCamelCase      ):
        if i == 0:
         UpperCAmelCase__							 =						0
        else:
         UpperCAmelCase__							 =						support_sizes[i - 1]

        UpperCAmelCase__							 =						S[s : s + size][start_token_masks[s : s + size]]
        UpperCAmelCase__							 =						S[s : s + size][end_token_masks[s : s + size]]

        UpperCAmelCase__							 =						torch.matmul(q[i] ,      s_start.T      ).sum(1      ).softmax(0      )
        UpperCAmelCase__							 =						torch.matmul(q[i] ,      s_end.T      ).sum(1      ).softmax(0      )

        if p_starts is not None:
         UpperCAmelCase__							 =						torch.vstack((p_starts, p_start)      )
         UpperCAmelCase__							 =						torch.vstack((p_ends, p_end)      )
        else:
         UpperCAmelCase__							 =						p_start
         UpperCAmelCase__							 =						p_end

       return p_starts, p_ends





