


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 		_SCREAMING_SNAKE_CASE     :


     def __init__(  self							,					A_							,					A_=13							,					A_=7							,					A_=6							,					A_=17							,					A_=23							,					A_=11							,					A_=True							,					):
       _UpperCAmelCase		:							List[Any]							=     parent
       _UpperCAmelCase		:							Optional[Any]							=     batch_size
       _UpperCAmelCase		:							Tuple							=     seq_length
       _UpperCAmelCase		:							List[Any]							=     act_dim
       _UpperCAmelCase		:							Any							=     state_dim
       _UpperCAmelCase		:							List[Any]							=     hidden_size
       _UpperCAmelCase		:							Union[str, Any]							=     max_length
       _UpperCAmelCase		:							Dict							=     is_training


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							Union[str, Any]							=     floats_tensor((self.batch_size, self.seq_length, self.state_dim)    )
       _UpperCAmelCase		:							Dict							=     floats_tensor((self.batch_size, self.seq_length, self.act_dim)    )
       _UpperCAmelCase		:							List[Any]							=     floats_tensor((self.batch_size, self.seq_length, 1)    )
       _UpperCAmelCase		:							Optional[int]							=     floats_tensor((self.batch_size, self.seq_length, 1)    )
       _UpperCAmelCase		:							Dict							=     ids_tensor((self.batch_size, self.seq_length)							,					vocab_size=10_00    )
       _UpperCAmelCase		:							Dict							=     random_attention_mask((self.batch_size, self.seq_length)    )

       _UpperCAmelCase		:							Union[str, Any]							=     self.get_config()

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


     def        __snake_case(  self    ):
       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        __snake_case(  self							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					):
       _UpperCAmelCase		:							Tuple							=     DecisionTransformerModel(config=A_    )
       model.to(A_    )
       model.eval()
       _UpperCAmelCase		:							Optional[Any]							=     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        __snake_case(  self    ):
       _UpperCAmelCase		:							List[Any]							=     self.prepare_config_and_inputs()
       (
           (
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),
       )		:							List[str]							=     config_and_inputs
       _UpperCAmelCase		:							Union[str, Any]							=     {
           """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 		_SCREAMING_SNAKE_CASE     (						A							,	A							,	A							,	unittest.TestCase     ):
     __SCREAMING_SNAKE_CASE     					=    (DecisionTransformerModel,) if is_torch_available() else ()
     __SCREAMING_SNAKE_CASE     					=    ()
     __SCREAMING_SNAKE_CASE     					=    {'''feature-extraction''': DecisionTransformerModel} if is_torch_available() else {}

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

     # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							Any							=     DecisionTransformerModelTester(self    )
       _UpperCAmelCase		:							Optional[Any]							=     ConfigTester(self							,					config_class=A_							,					hidden_size=37    )


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


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							List[Any]							=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_    )


     @slow
     def        __snake_case(  self    ):
       for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _UpperCAmelCase		:							List[str]							=     DecisionTransformerModel.from_pretrained(A_    )
         self.assertIsNotNone(A_    )




     def        __snake_case(  self    ):
       _UpperCAmelCase,_UpperCAmelCase		:							Optional[Any]							=     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         _UpperCAmelCase		:							Optional[Any]							=     model_class(A_    )
         _UpperCAmelCase		:							int							=     inspect.signature(model.forward    )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         _UpperCAmelCase		:							Union[str, Any]							=     [*signature.parameters.keys()]

         _UpperCAmelCase		:							int							=     [
             """states""",
             """actions""",
             """rewards""",
             """returns_to_go""",
             """timesteps""",
             """attention_mask""",
         ]

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


@require_torch
class 		_SCREAMING_SNAKE_CASE     (						unittest.TestCase     ):


     @slow
     def        __snake_case(  self    ):
       _UpperCAmelCase		:							Any							=     2  # number of steps of autoregressive prediction we will perform
       _UpperCAmelCase		:							Dict							=     10  # defined by the RL environment, may be normalized
       _UpperCAmelCase		:							Union[str, Any]							=     DecisionTransformerModel.from_pretrained("""edbeeching/decision-transformer-gym-hopper-expert"""    )
       _UpperCAmelCase		:							str							=     model.to(A_    )
       _UpperCAmelCase		:							str							=     model.config
       torch.manual_seed(0    )
       _UpperCAmelCase		:							List[Any]							=     torch.randn(1							,					1							,					config.state_dim    ).to(device=A_							,					dtype=torch.floataa    )  # env.reset()

       _UpperCAmelCase		:							int							=     torch.tensor(
           [[0.2_4_2_7_9_3, -0.2_8_6_9_3_0_7_4, 0.8_7_4_2_6_1_3], [0.6_7_8_1_5_2_7_4, -0.0_8_1_0_1_0_8_5, -0.1_2_9_5_2_1_4_7]]							,					device=A_    )

       _UpperCAmelCase		:							int							=     torch.tensor(A_							,					device=A_							,					dtype=torch.floataa    ).reshape(1							,					1							,					1    )
       _UpperCAmelCase		:							Any							=     state
       _UpperCAmelCase		:							str							=     torch.zeros(1							,					0							,					config.act_dim							,					device=A_							,					dtype=torch.floataa    )
       _UpperCAmelCase		:							Union[str, Any]							=     torch.zeros(1							,					0							,					device=A_							,					dtype=torch.floataa    )
       _UpperCAmelCase		:							Union[str, Any]							=     torch.tensor(0							,					device=A_							,					dtype=torch.long    ).reshape(1							,					1    )

       for step in range(A_    ):
         _UpperCAmelCase		:							Tuple							=     torch.cat([actions, torch.zeros(1							,					1							,					config.act_dim							,					device=A_    )]							,					dim=1    )
         _UpperCAmelCase		:							Union[str, Any]							=     torch.cat([rewards, torch.zeros(1							,					1							,					device=A_    )]							,					dim=1    )

         _UpperCAmelCase		:							List[str]							=     torch.ones(1							,					states.shape[1]    ).to(dtype=torch.long							,					device=states.device    )

         with torch.no_grad():
           _UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase		:							Dict							=     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    )    )
         _UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase		:							str							=     (  # env.step(action)
             torch.randn(1							,					1							,					config.state_dim    ).to(device=A_							,					dtype=torch.floataa    ),
             1.0,
             False,
             {},
         )

         _UpperCAmelCase		:							Any							=     action_pred[0, -1]
         _UpperCAmelCase		:							str							=     torch.cat([states, state]							,					dim=1    )
         _UpperCAmelCase		:							List[Any]							=     returns_to_go[0, -1] - reward
         _UpperCAmelCase		:							Union[str, Any]							=     torch.cat([returns_to_go, pred_return.reshape(1							,					1							,					1    )]							,					dim=1    )
         _UpperCAmelCase		:							int							=     torch.cat(
             [timesteps, torch.ones((1, 1)							,					device=A_							,					dtype=torch.long    ) * (step + 1)]							,					dim=1    )









import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE__					:     Optional[Any]       =     logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__					:     Tuple       =     {
    'b0': efficientnet.EfficientNetBa,
    'b1': efficientnet.EfficientNetBa,
    'b2': efficientnet.EfficientNetBa,
    'b3': efficientnet.EfficientNetBa,
    'b4': efficientnet.EfficientNetBa,
    'b5': efficientnet.EfficientNetBa,
    'b6': efficientnet.EfficientNetBa,
    'b7': efficientnet.EfficientNetBa,
}

SCREAMING_SNAKE_CASE__					:     Optional[int]       =     {
    'b0': {
        'hidden_dim': 1_280,
        'width_coef': 1.0,
        'depth_coef': 1.0,
        'image_size': 224,
        'dropout_rate': 0.2,
        'dw_padding': [],
    },
    'b1': {
        'hidden_dim': 1_280,
        'width_coef': 1.0,
        'depth_coef': 1.1,
        'image_size': 240,
        'dropout_rate': 0.2,
        'dw_padding': [16],
    },
    'b2': {
        'hidden_dim': 1_408,
        'width_coef': 1.1,
        'depth_coef': 1.2,
        'image_size': 260,
        'dropout_rate': 0.3,
        'dw_padding': [5, 8, 16],
    },
    'b3': {
        'hidden_dim': 1_536,
        'width_coef': 1.2,
        'depth_coef': 1.4,
        'image_size': 300,
        'dropout_rate': 0.3,
        'dw_padding': [5, 18],
    },
    'b4': {
        'hidden_dim': 1_792,
        'width_coef': 1.4,
        'depth_coef': 1.8,
        'image_size': 380,
        'dropout_rate': 0.4,
        'dw_padding': [6],
    },
    'b5': {
        'hidden_dim': 2_048,
        'width_coef': 1.6,
        'depth_coef': 2.2,
        'image_size': 456,
        'dropout_rate': 0.4,
        'dw_padding': [13, 27],
    },
    'b6': {
        'hidden_dim': 2_304,
        'width_coef': 1.8,
        'depth_coef': 2.6,
        'image_size': 528,
        'dropout_rate': 0.5,
        'dw_padding': [31],
    },
    'b7': {
        'hidden_dim': 2_560,
        'width_coef': 2.0,
        'depth_coef': 3.1,
        'image_size': 600,
        'dropout_rate': 0.5,
        'dw_padding': [18],
    },
}




def    a__    (		snake_case__					:							Optional[Any]     ):
  _UpperCAmelCase		:							Optional[int]							=     EfficientNetConfig()
  _UpperCAmelCase		:							Tuple							=     CONFIG_MAP[model_name]["""hidden_dim"""]
  _UpperCAmelCase		:							Any							=     CONFIG_MAP[model_name]["""width_coef"""]
  _UpperCAmelCase		:							Optional[Any]							=     CONFIG_MAP[model_name]["""depth_coef"""]
  _UpperCAmelCase		:							int							=     CONFIG_MAP[model_name]["""image_size"""]
  _UpperCAmelCase		:							List[str]							=     CONFIG_MAP[model_name]["""dropout_rate"""]
  _UpperCAmelCase		:							Optional[int]							=     CONFIG_MAP[model_name]["""dw_padding"""]

  _UpperCAmelCase		:							int							=     """huggingface/label-files"""
  _UpperCAmelCase		:							Dict							=     """imagenet-1k-id2label.json"""
  _UpperCAmelCase		:							Optional[Any]							=     1000
  _UpperCAmelCase		:							List[str]							=     json.load(open(hf_hub_download(snake_case__						,   snake_case__						,   repo_type="""dataset"""     )						,   """r"""     )     )
  _UpperCAmelCase		:							Union[str, Any]							=     {int(snake_case__     ): v for k, v in idalabel.items()}

  _UpperCAmelCase		:							Any							=     idalabel
  _UpperCAmelCase		:							List[str]							=     {v: k for k, v in idalabel.items()}
  return config




def    a__    (		):
  _UpperCAmelCase		:							Optional[int]							=     """http://images.cocodataset.org/val2017/000000039769.jpg"""
  _UpperCAmelCase		:							str							=     Image.open(requests.get(snake_case__						,   stream=snake_case__     ).raw     )
  return im




def    a__    (		snake_case__					:							Tuple     ):
  _UpperCAmelCase		:							List[str]							=     CONFIG_MAP[model_name]["""image_size"""]
  _UpperCAmelCase		:							str							=     EfficientNetImageProcessor(
      size={"""height""": size, """width""": size}						,   image_mean=[0.485, 0.456, 0.406]						,   image_std=[0.4785_3944, 0.473_2864, 0.4743_4163]						,   do_center_crop=snake_case__						,   )
  return preprocessor




def    a__    (		snake_case__					:							Optional[Any]     ):
  _UpperCAmelCase		:							Optional[int]							=     [v.split("""_"""     )[0].split("""block"""     )[1] for v in original_param_names if v.startswith("""block"""     )]
  _UpperCAmelCase		:							Dict							=     sorted(set(snake_case__     )     )
  _UpperCAmelCase		:							int							=     len(snake_case__     )
  _UpperCAmelCase		:							Optional[Any]							=     {b: str(snake_case__     ) for b, i in zip(snake_case__						,   range(snake_case__     )     )}

  _UpperCAmelCase		:							Union[str, Any]							=     []
  rename_keys.append(("""stem_conv/kernel:0""", """embeddings.convolution.weight""")     )
  rename_keys.append(("""stem_bn/gamma:0""", """embeddings.batchnorm.weight""")     )
  rename_keys.append(("""stem_bn/beta:0""", """embeddings.batchnorm.bias""")     )
  rename_keys.append(("""stem_bn/moving_mean:0""", """embeddings.batchnorm.running_mean""")     )
  rename_keys.append(("""stem_bn/moving_variance:0""", """embeddings.batchnorm.running_var""")     )

  for b in block_names:
    _UpperCAmelCase		:							List[Any]							=     block_name_mapping[b]
    rename_keys.append((f'''block{b}_expand_conv/kernel:0''', f'''encoder.blocks.{hf_b}.expansion.expand_conv.weight''')     )
    rename_keys.append((f'''block{b}_expand_bn/gamma:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.weight''')     )
    rename_keys.append((f'''block{b}_expand_bn/beta:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.bias''')     )
    rename_keys.append(
        (f'''block{b}_expand_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.running_mean''')     )
    rename_keys.append(
        (f'''block{b}_expand_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.expansion.expand_bn.running_var''')     )
    rename_keys.append(
        (f'''block{b}_dwconv/depthwise_kernel:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight''')     )
    rename_keys.append((f'''block{b}_bn/gamma:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight''')     )
    rename_keys.append((f'''block{b}_bn/beta:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias''')     )
    rename_keys.append(
        (f'''block{b}_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean''')     )
    rename_keys.append(
        (f'''block{b}_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var''')     )

    rename_keys.append((f'''block{b}_se_reduce/kernel:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.reduce.weight''')     )
    rename_keys.append((f'''block{b}_se_reduce/bias:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.reduce.bias''')     )
    rename_keys.append((f'''block{b}_se_expand/kernel:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.expand.weight''')     )
    rename_keys.append((f'''block{b}_se_expand/bias:0''', f'''encoder.blocks.{hf_b}.squeeze_excite.expand.bias''')     )
    rename_keys.append(
        (f'''block{b}_project_conv/kernel:0''', f'''encoder.blocks.{hf_b}.projection.project_conv.weight''')     )
    rename_keys.append((f'''block{b}_project_bn/gamma:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.weight''')     )
    rename_keys.append((f'''block{b}_project_bn/beta:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.bias''')     )
    rename_keys.append(
        (f'''block{b}_project_bn/moving_mean:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.running_mean''')     )
    rename_keys.append(
        (f'''block{b}_project_bn/moving_variance:0''', f'''encoder.blocks.{hf_b}.projection.project_bn.running_var''')     )

  rename_keys.append(("""top_conv/kernel:0""", """encoder.top_conv.weight""")     )
  rename_keys.append(("""top_bn/gamma:0""", """encoder.top_bn.weight""")     )
  rename_keys.append(("""top_bn/beta:0""", """encoder.top_bn.bias""")     )
  rename_keys.append(("""top_bn/moving_mean:0""", """encoder.top_bn.running_mean""")     )
  rename_keys.append(("""top_bn/moving_variance:0""", """encoder.top_bn.running_var""")     )

  _UpperCAmelCase		:							Union[str, Any]							=     {}
  for item in rename_keys:
    if item[0] in original_param_names:
      _UpperCAmelCase		:							Tuple							=     """efficientnet.""" + item[1]

  _UpperCAmelCase		:							List[Any]							=     """classifier.weight"""
  _UpperCAmelCase		:							List[Any]							=     """classifier.bias"""
  return key_mapping




def    a__    (		snake_case__					:							Optional[int]						,   snake_case__					:							int						,   snake_case__					:							List[str]     ):
  for key, value in tf_params.items():
    if "normalization" in key:
      continue

    _UpperCAmelCase		:							Optional[Any]							=     key_mapping[key]
    if "_conv" in key and "kernel" in key:
      _UpperCAmelCase		:							Optional[Any]							=     torch.from_numpy(snake_case__     ).permute(3						,   2						,   0						,   1     )
    elif "depthwise_kernel" in key:
      _UpperCAmelCase		:							Dict							=     torch.from_numpy(snake_case__     ).permute(2						,   3						,   0						,   1     )
    elif "kernel" in key:
      _UpperCAmelCase		:							Any							=     torch.from_numpy(np.transpose(snake_case__     )     )
    else:
      _UpperCAmelCase		:							Optional[Any]							=     torch.from_numpy(snake_case__     )

    # Replace HF parameters with original TF model parameters
    assert hf_params[hf_key].shape == new_hf_value.shape
    hf_params[hf_key].copy_(snake_case__     )




@torch.no_grad()
def    a__    (		snake_case__					:							str						,   snake_case__					:							Dict						,   snake_case__					:							int						,   snake_case__					:							str     ):
  _UpperCAmelCase		:							List[str]							=     model_classes[model_name](
      include_top=snake_case__						,   weights="""imagenet"""						,   input_tensor=snake_case__						,   input_shape=snake_case__						,   pooling=snake_case__						,   classes=1000						,   classifier_activation="""softmax"""						,   )

  _UpperCAmelCase		:							Any							=     original_model.trainable_variables
  _UpperCAmelCase		:							List[Any]							=     original_model.non_trainable_variables
  _UpperCAmelCase		:							List[Any]							=     {param.name: param.numpy() for param in tf_params}
  for param in tf_non_train_params:
    _UpperCAmelCase		:							List[str]							=     param.numpy()
  _UpperCAmelCase		:							List[Any]							=     list(tf_params.keys()     )

  # Load HuggingFace model
  _UpperCAmelCase		:							Any							=     get_efficientnet_config(snake_case__     )
  _UpperCAmelCase		:							List[str]							=     EfficientNetForImageClassification(snake_case__     ).eval()
  _UpperCAmelCase		:							Any							=     hf_model.state_dict()

  # Create src-to-dst parameter name mapping dictionary
  print("""Converting parameters..."""     )
  _UpperCAmelCase		:							Optional[Any]							=     rename_keys(snake_case__     )
  replace_params(snake_case__						,   snake_case__						,   snake_case__     )

  # Initialize preprocessor and preprocess input image
  _UpperCAmelCase		:							Dict							=     convert_image_processor(snake_case__     )
  _UpperCAmelCase		:							List[Any]							=     preprocessor(images=prepare_img()						,   return_tensors="""pt"""     )

  # HF model inference
  hf_model.eval()
  with torch.no_grad():
    _UpperCAmelCase		:							List[str]							=     hf_model(**snake_case__     )
  _UpperCAmelCase		:							Any							=     outputs.logits.detach().numpy()

  # Original model inference
  _UpperCAmelCase		:							Dict							=     False
  _UpperCAmelCase		:							List[Any]							=     CONFIG_MAP[model_name]["""image_size"""]
  _UpperCAmelCase		:							Any							=     prepare_img().resize((image_size, image_size)						,   resample=PIL.Image.NEAREST     )
  _UpperCAmelCase		:							Optional[Any]							=     image.img_to_array(snake_case__     )
  _UpperCAmelCase		:							Any							=     np.expand_dims(snake_case__						,   axis=0     )
  _UpperCAmelCase		:							List[Any]							=     original_model.predict(snake_case__     )

  # Check whether original and HF model outputs match  -> np.allclose
  assert np.allclose(snake_case__						,   snake_case__						,   atol=1e-3     ), "The predicted logits are not the same."
  print("""Model outputs match!"""     )

  if save_model:
    # Create folder to save model
    if not os.path.isdir(snake_case__     ):
      os.mkdir(snake_case__     )
    # Save converted model and image processor
    hf_model.save_pretrained(snake_case__     )
    preprocessor.save_pretrained(snake_case__     )

  if push_to_hub:
    # Push model and image processor to hub
    print(f'''Pushing converted {model_name} to the hub...'''     )
    _UpperCAmelCase		:							int							=     f'''efficientnet-{model_name}'''
    preprocessor.push_to_hub(snake_case__     )
    hf_model.push_to_hub(snake_case__     )


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE__					:     List[Any]       =     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--model_name',
     default='b0',
     type=str,
     help='Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].',
 )
 parser.add_argument(
     '--pytorch_dump_folder_path',
     default='hf_model',
     type=str,
     help='Path to the output PyTorch model directory.',
 )
 parser.add_argument('--save_model', action='store_true', help='Save model to local')
 parser.add_argument('--push_to_hub', action='store_true', help='Push model and image processor to the hub')

 SCREAMING_SNAKE_CASE__					:     Optional[Any]       =     parser.parse_args()
 convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)





