


import tempfile
import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import DPRContextEncoder, DPRQuestionEncoder, DPRReader, DPRReaderTokenizer
   from transformers.models.dpr.modeling_dpr import (
       DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
       DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
       DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
   )







class lowerCAmelCase_						:

 '''simple docstring'''

 def __init__(	self		,      _UpperCAmelCase		,      _UpperCAmelCase=13		,      _UpperCAmelCase=7		,      _UpperCAmelCase=False		,      _UpperCAmelCase=True		,      _UpperCAmelCase=True		,      _UpperCAmelCase=True		,      _UpperCAmelCase=99		,      _UpperCAmelCase=32		,      _UpperCAmelCase=5		,      _UpperCAmelCase=4		,      _UpperCAmelCase=37		,      _UpperCAmelCase="gelu"		,      _UpperCAmelCase=0.1		,      _UpperCAmelCase=0.1		,      _UpperCAmelCase=5_12		,      _UpperCAmelCase=16		,      _UpperCAmelCase=2		,      _UpperCAmelCase=0.02		,      _UpperCAmelCase=3		,      _UpperCAmelCase=4		,      _UpperCAmelCase=None		,      _UpperCAmelCase=0		,      ):
       snake_case_       =      parent
       snake_case_       =      batch_size
       snake_case_       =      seq_length
       snake_case_       =      is_training
       snake_case_       =      use_input_mask
       snake_case_       =      use_token_type_ids
       snake_case_       =      use_labels
       snake_case_       =      vocab_size
       snake_case_       =      hidden_size
       snake_case_       =      num_hidden_layers
       snake_case_       =      num_attention_heads
       snake_case_       =      intermediate_size
       snake_case_       =      hidden_act
       snake_case_       =      hidden_dropout_prob
       snake_case_       =      attention_probs_dropout_prob
       snake_case_       =      max_position_embeddings
       snake_case_       =      type_vocab_size
       snake_case_       =      type_sequence_label_size
       snake_case_       =      initializer_range
       snake_case_       =      num_labels
       snake_case_       =      num_choices
       snake_case_       =      scope
       snake_case_       =      projection_dim



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size					)

       snake_case_       =      None
       if self.use_input_mask:
             snake_case_       =      random_attention_mask([self.batch_size, self.seq_length]					)

       snake_case_       =      None
       if self.use_token_type_ids:
             snake_case_       =      ids_tensor([self.batch_size, self.seq_length]		,      self.type_vocab_size					)

       snake_case_       =      None
       snake_case_       =      None
       snake_case_       =      None
       if self.use_labels:
             snake_case_       =      ids_tensor([self.batch_size]		,      self.type_sequence_label_size					)
             snake_case_       =      ids_tensor([self.batch_size, self.seq_length]		,      self.num_labels					)
             snake_case_       =      ids_tensor([self.batch_size]		,      self.num_choices					)

       snake_case_       =      self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



 def 						UpperCamelCase__     (	self					):
       return DPRConfig(
           projection_dim=self.projection_dim		,      vocab_size=self.vocab_size		,      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		,      max_position_embeddings=self.max_position_embeddings		,      type_vocab_size=self.type_vocab_size		,      initializer_range=self.initializer_range		,      )



 def 						UpperCamelCase__     (	self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase					):
       snake_case_       =      DPRContextEncoder(config=_UpperCAmelCase					)
       model.to(_UpperCAmelCase					)
       model.eval()
       snake_case_       =      model(_UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      token_type_ids=_UpperCAmelCase					)
       snake_case_       =      model(_UpperCAmelCase		,      token_type_ids=_UpperCAmelCase					)
       snake_case_       =      model(_UpperCAmelCase					)
       self.parent.assertEqual(result.pooler_output.shape		,      (self.batch_size, self.projection_dim or self.hidden_size)					)



 def 						UpperCamelCase__     (	self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase					):
       snake_case_       =      DPRQuestionEncoder(config=_UpperCAmelCase					)
       model.to(_UpperCAmelCase					)
       model.eval()
       snake_case_       =      model(_UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      token_type_ids=_UpperCAmelCase					)
       snake_case_       =      model(_UpperCAmelCase		,      token_type_ids=_UpperCAmelCase					)
       snake_case_       =      model(_UpperCAmelCase					)
       self.parent.assertEqual(result.pooler_output.shape		,      (self.batch_size, self.projection_dim or self.hidden_size)					)



 def 						UpperCamelCase__     (	self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase					):
       snake_case_       =      DPRReader(config=_UpperCAmelCase					)
       model.to(_UpperCAmelCase					)
       model.eval()
       snake_case_       =      model(
           _UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      )

       self.parent.assertEqual(result.start_logits.shape		,      (self.batch_size, self.seq_length)					)
       self.parent.assertEqual(result.end_logits.shape		,      (self.batch_size, self.seq_length)					)
       self.parent.assertEqual(result.relevance_logits.shape		,      (self.batch_size,)					)





 def 						UpperCamelCase__     (	self					):
       snake_case_       =      self.prepare_config_and_inputs()
       (
           (
           snake_case_
       )       ,    (
           snake_case_
       )       ,    (
           snake_case_
       )       ,    (
           snake_case_
       )       ,    (
           snake_case_
       )       ,    (
           snake_case_
       )       ,    (
           snake_case_
       )       ,    
       )       =      config_and_inputs
       snake_case_       =      {'''input_ids''': input_ids}
       return config, inputs_dict







@require_torch
class lowerCAmelCase_						( lowerCamelCase__					,						lowerCamelCase__					,						unittest.TestCase       ):

 '''simple docstring'''

 __snake_case												=  (
     (
         DPRContextEncoder,
         DPRQuestionEncoder,
         DPRReader,
     )
     if is_torch_available()
     else ()
 )
 __snake_case												=  {"feature-extraction": DPRQuestionEncoder} if is_torch_available() else {}

 __snake_case												=  False
 __snake_case												=  False  # why?
 __snake_case												=  False
 __snake_case												=  False



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      DPRModelTester(self					)
       snake_case_       =      ConfigTester(self		,      config_class=_UpperCAmelCase		,      hidden_size=37					)



 def 						UpperCamelCase__     (	self					):
       self.config_tester.run_common_tests()



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_context_encoder(*_UpperCAmelCase					)



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_question_encoder(*_UpperCAmelCase					)



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_reader(*_UpperCAmelCase					)



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      self.model_tester.prepare_config_and_inputs()[0]

       snake_case_       =      DPRQuestionEncoder(config=_UpperCAmelCase					)
       model.to(_UpperCAmelCase					)
       model.eval()

       with tempfile.TemporaryDirectory() as tmp_dirname:
             model.save_pretrained(_UpperCAmelCase					)
             snake_case_       =      DPRQuestionEncoder.from_pretrained(_UpperCAmelCase		,      projection_dim=5_12					)

       self.assertIsNotNone(_UpperCAmelCase					)





 @slow
 def 						UpperCamelCase__     (	self					):
       for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_       =      DPRContextEncoder.from_pretrained(_UpperCAmelCase					)
             self.assertIsNotNone(_UpperCAmelCase					)

       for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_       =      DPRContextEncoder.from_pretrained(_UpperCAmelCase					)
             self.assertIsNotNone(_UpperCAmelCase					)

       for model_name in DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_       =      DPRQuestionEncoder.from_pretrained(_UpperCAmelCase					)
             self.assertIsNotNone(_UpperCAmelCase					)

       for model_name in DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_       =      DPRReader.from_pretrained(_UpperCAmelCase					)
             self.assertIsNotNone(_UpperCAmelCase					)







@require_torch
class lowerCAmelCase_						( unittest.TestCase       ):

 '''simple docstring'''

 @slow
 def 						UpperCamelCase__     (	self					):
       snake_case_       =      DPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base'''		,      return_dict=_UpperCAmelCase					)
       model.to(_UpperCAmelCase					)

       snake_case_       =      torch.tensor(
           [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]]		,      dtype=torch.long		,      device=_UpperCAmelCase					)  # [CLS] hello, is my dog cute? [SEP]
       snake_case_       =      model(_UpperCAmelCase					)[0]  # embedding shape = (1, 768)
       # compare the actual values for a slice.
       snake_case_       =      torch.tensor(
           [
               [
                   0.03_236_253,
                   0.12_753_335,
                   0.16_818_509,
                   0.00_279_786,
                   0.3_896_933,
                   0.24_264_945,
                   0.2_178_971,
                   -0.02_335_227,
                   -0.08_481_959,
                   -0.14_324_117,
               ]
           ]		,      dtype=torch.float		,      device=_UpperCAmelCase		,      )
       self.assertTrue(torch.allclose(output[:, :10]		,      _UpperCAmelCase		,      atol=1E-4					)					)





 @slow
 def 						UpperCamelCase__     (	self					):
       snake_case_       =      DPRReaderTokenizer.from_pretrained('''facebook/dpr-reader-single-nq-base'''					)
       snake_case_       =      DPRReader.from_pretrained('''facebook/dpr-reader-single-nq-base'''					)
       model.to(_UpperCAmelCase					)

       snake_case_       =      tokenizer(
           questions='''What is love ?'''		,      titles='''Haddaway'''		,      texts='''What Is Love is a song recorded by the artist Haddaway'''		,      padding=_UpperCAmelCase		,      return_tensors='''pt'''		,      )
       encoded_inputs.to(_UpperCAmelCase					)

       snake_case_       =      model(**_UpperCAmelCase					)

       # compare the actual values for a slice.
       snake_case_       =      torch.tensor(
           [[-10.3_005, -10.7_765, -11.4_872, -11.6_841, -11.9_312, -10.3_002, -9.8_544, -11.7_378, -12.0_821, -10.2_975]]		,      dtype=torch.float		,      device=_UpperCAmelCase		,      )

       snake_case_       =      torch.tensor(
           [[-11.0_684, -11.7_041, -11.5_397, -10.3_465, -10.8_791, -6.8_443, -11.9_959, -11.0_364, -10.0_096, -6.8_405]]		,      dtype=torch.float		,      device=_UpperCAmelCase		,      )
       self.assertTrue(torch.allclose(outputs.start_logits[:, :10]		,      _UpperCAmelCase		,      atol=1E-4					)					)
       self.assertTrue(torch.allclose(outputs.end_logits[:, :10]		,      _UpperCAmelCase		,      atol=1E-4					)					)



import argparse
import json
import re
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileNetVaConfig,
    MobileNetVaForImageClassification,
    MobileNetVaForSemanticSegmentation,
    MobileNetVaImageProcessor,
    load_tf_weights_in_mobilenet_va,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase			=    logging.get_logger(__name__)






def        __lowerCAmelCase    (SCREAMING_SNAKE_CASE					)->  Any:




      """simple docstring"""

      snake_case_       =      MobileNetVaConfig(layer_norm_eps=0.0_0_1					)

      if "quant" in model_name:
            raise ValueError('''Quantized models are not supported.'''					)

      snake_case_       =      re.match(R'''^.*mobilenet_v2_([^_]*)_([^_]*)$'''					,					SCREAMING_SNAKE_CASE					)
      if matches:
            snake_case_       =      float(matches[1]					)
            snake_case_       =      int(matches[2]					)

      if model_name.startswith('''deeplabv3_'''					):
            snake_case_       =      8
            snake_case_       =      21
            snake_case_       =      '''pascal-voc-id2label.json'''
      else:
            # The TensorFlow version of MobileNetV2 predicts 1001 classes instead
            # of the usual 1000. The first class (index 0) is "background".
            snake_case_       =      1001
            snake_case_       =      '''imagenet-1k-id2label.json'''

      snake_case_       =      '''huggingface/label-files'''
      snake_case_       =      json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE					,					repo_type='''dataset'''					)					,					'''r'''					)					)

      if config.num_labels == 1001:
            snake_case_       =      {int(SCREAMING_SNAKE_CASE					) + 1: v for k, v in idalabel.items()}
            snake_case_       =      '''background'''
      else:
            snake_case_       =      {int(SCREAMING_SNAKE_CASE					): v for k, v in idalabel.items()}

      snake_case_       =      idalabel
      snake_case_       =      {v: k for k, v in idalabel.items()}

      return config






def        __lowerCAmelCase    ()->  str:




      """simple docstring"""

      snake_case_       =      '''http://images.cocodataset.org/val2017/000000039769.jpg'''
      snake_case_       =      Image.open(requests.get(SCREAMING_SNAKE_CASE					,					stream=SCREAMING_SNAKE_CASE					).raw					)
      return im






@torch.no_grad()
def        __lowerCAmelCase    (SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE=False					)->  str:




      """simple docstring"""

      snake_case_       =      get_mobilenet_va_config(SCREAMING_SNAKE_CASE					)

      # Load 🤗 model
      if model_name.startswith('''deeplabv3_'''					):
            snake_case_       =      MobileNetVaForSemanticSegmentation(SCREAMING_SNAKE_CASE					).eval()
      else:
            snake_case_       =      MobileNetVaForImageClassification(SCREAMING_SNAKE_CASE					).eval()

      # Load weights from TensorFlow checkpoint
      load_tf_weights_in_mobilenet_va(SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE					)

      # Check outputs on an image, prepared by MobileNetV2ImageProcessor
      snake_case_       =      MobileNetVaImageProcessor(
          crop_size={'''width''': config.image_size, '''height''': config.image_size}					,					size={'''shortest_edge''': config.image_size + 32}					,					)
      snake_case_       =      image_processor(images=prepare_img()					,					return_tensors='''pt'''					)
      snake_case_       =      model(**SCREAMING_SNAKE_CASE					)
      snake_case_       =      outputs.logits

      if model_name.startswith('''deeplabv3_'''					):
            assert logits.shape == (1, 21, 65, 65)

            if model_name == "deeplabv3_mobilenet_v2_1.0_513":
                  snake_case_       =      torch.tensor(
                      [
                          [[1_7.5_7_9_0, 1_7.7_5_8_1, 1_8.3_3_5_5], [1_8.3_2_5_7, 1_8.4_2_3_0, 1_8.8_9_7_3], [1_8.6_1_6_9, 1_8.8_6_5_0, 1_9.2_1_8_7]],
                          [[-2.1_5_9_5, -2.0_9_7_7, -2.3_7_4_1], [-2.4_2_2_6, -2.3_0_2_8, -2.6_8_3_5], [-2.7_8_1_9, -2.5_9_9_1, -2.7_7_0_6]],
                          [[4.2_0_5_8, 4.8_3_1_7, 4.7_6_3_8], [4.4_1_3_6, 5.0_3_6_1, 4.9_3_8_3], [4.5_0_2_8, 4.9_6_4_4, 4.8_7_3_4]],
                      ]					)

            else:
                  raise ValueError(f'''Unknown model name: {model_name}'''					)

            assert torch.allclose(logits[0, :3, :3, :3]					,					SCREAMING_SNAKE_CASE					,					atol=1E-4					)
      else:
            assert logits.shape == (1, 1001)

            if model_name == "mobilenet_v2_1.4_224":
                  snake_case_       =      torch.tensor([0.0_1_8_1, -1.0_0_1_5, 0.4_6_8_8]					)
            elif model_name == "mobilenet_v2_1.0_224":
                  snake_case_       =      torch.tensor([0.2_4_4_5, -1.1_9_9_3, 0.1_9_0_5]					)
            elif model_name == "mobilenet_v2_0.75_160":
                  snake_case_       =      torch.tensor([0.2_4_8_2, 0.4_1_3_6, 0.6_6_6_9]					)
            elif model_name == "mobilenet_v2_0.35_96":
                  snake_case_       =      torch.tensor([0.1_4_5_1, -0.4_6_2_4, 0.7_1_9_2]					)
            else:
                  snake_case_       =      None

            if expected_logits is not None:
                  assert torch.allclose(logits[0, :3]					,					SCREAMING_SNAKE_CASE					,					atol=1E-4					)

      Path(SCREAMING_SNAKE_CASE					).mkdir(exist_ok=SCREAMING_SNAKE_CASE					)
      print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''					)
      model.save_pretrained(SCREAMING_SNAKE_CASE					)
      print(f'''Saving image processor to {pytorch_dump_folder_path}'''					)
      image_processor.save_pretrained(SCREAMING_SNAKE_CASE					)

      if push_to_hub:
            print('''Pushing to the hub...'''					)
            snake_case_       =      '''google/''' + model_name
            image_processor.push_to_hub(SCREAMING_SNAKE_CASE					)
            model.push_to_hub(SCREAMING_SNAKE_CASE					)


if __name__ == "__main__":
   UpperCAmelCase			=    argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--model_name""",
       default="""mobilenet_v2_1.0_224""",
       type=str,
       help="""Name of the MobileNetV2 model you'd like to convert. Should in the form 'mobilenet_v2_<depth>_<size>'.""",
   )
   parser.add_argument(
       """--checkpoint_path""", required=True, type=str, help="""Path to the original TensorFlow checkpoint (.ckpt file)."""
   )
   parser.add_argument(
       """--pytorch_dump_folder_path""", required=True, type=str, help="""Path to the output PyTorch model directory."""
   )
   parser.add_argument(
       """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
   )

   UpperCAmelCase			=    parser.parse_args()
   convert_movilevit_checkpoint(
       args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
   )