

"""simple docstring"""




from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import TFCvtForImageClassification, TFCvtModel
      from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor







class   __lowercase      (     __lowerCamelCase       ):

  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Union[str, Any]        =       self.config_class(**self.inputs_dict     )
   self.parent.assertTrue(hasattr(A       ,"""embed_dim"""     )     )
   self.parent.assertTrue(hasattr(A       ,"""num_heads"""     )     )







class   __lowercase      :

  def __init__(	self		:      Optional[int]       ,A		:      Tuple       ,A		:      Any=13       ,A		:      str=64       ,A		:      List[str]=3       ,A		:      int=[16, 48, 96]       ,A		:      Optional[Any]=[1, 3, 6]       ,A		:      Any=[1, 2, 10]       ,A		:      List[str]=[7, 3, 3]       ,A		:      str=[4, 2, 2]       ,A		:      str=[2, 1, 1]       ,A		:      Union[str, Any]=[2, 2, 2]       ,A		:      List[str]=[False, False, True]       ,A		:      Dict=[0.0, 0.0, 0.0]       ,A		:      Tuple=0.0_2       ,A		:      List[Any]=1e-12       ,A		:      Dict=True       ,A		:      Optional[Any]=True       ,A		:      Optional[int]=2       ,):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       parent
   UpperCAmelCase__       :	List[str]        =       batch_size
   UpperCAmelCase__       :	int        =       image_size
   UpperCAmelCase__       :	List[str]        =       patch_sizes
   UpperCAmelCase__       :	int        =       patch_stride
   UpperCAmelCase__       :	Tuple        =       patch_padding
   UpperCAmelCase__       :	Tuple        =       is_training
   UpperCAmelCase__       :	Dict        =       use_labels
   UpperCAmelCase__       :	Dict        =       num_labels
   UpperCAmelCase__       :	Tuple        =       num_channels
   UpperCAmelCase__       :	Any        =       embed_dim
   UpperCAmelCase__       :	Tuple        =       num_heads
   UpperCAmelCase__       :	Union[str, Any]        =       stride_kv
   UpperCAmelCase__       :	List[str]        =       depth
   UpperCAmelCase__       :	Optional[Any]        =       cls_token
   UpperCAmelCase__       :	List[Any]        =       attention_drop_rate
   UpperCAmelCase__       :	Any        =       initializer_range
   UpperCAmelCase__       :	Optional[int]        =       layer_norm_eps

  def  __lowercase							(	self		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	int        =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

   UpperCAmelCase__       :	Any        =       None
   if self.use_labels:
    # create a random int32 tensor of given shape
    UpperCAmelCase__       :	Union[str, Any]        =       ids_tensor([self.batch_size]       ,self.num_labels     )

   UpperCAmelCase__       :	Dict        =       self.get_config()
   return config, pixel_values, labels

  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   return CvtConfig(
       image_size=self.image_size       ,num_labels=self.num_labels       ,num_channels=self.num_channels       ,embed_dim=self.embed_dim       ,num_heads=self.num_heads       ,patch_sizes=self.patch_sizes       ,patch_padding=self.patch_padding       ,patch_stride=self.patch_stride       ,stride_kv=self.stride_kv       ,depth=self.depth       ,cls_token=self.cls_token       ,attention_drop_rate=self.attention_drop_rate       ,initializer_range=self.initializer_range       ,)

  def  __lowercase							(	self		:      List[str]       ,A		:      Any       ,A		:      List[Any]       ,A		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       TFCvtModel(config=A     )
   UpperCAmelCase__       :	Dict        =       model(A       ,training=A     )
   UpperCAmelCase__       :	List[str]        =       (self.image_size, self.image_size)
   UpperCAmelCase__ ,   UpperCAmelCase__       :	Tuple        =       image_size[0], image_size[1]
   for i in range(len(self.depth     )     ):
    UpperCAmelCase__       :	Optional[Any]        =       floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1     )
    UpperCAmelCase__       :	List[Any]        =       floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1     )
   self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.embed_dim[-1], height, width)     )

  def  __lowercase							(	self		:      int       ,A		:      List[Any]       ,A		:      int       ,A		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Union[str, Any]        =       self.num_labels
   UpperCAmelCase__       :	List[Any]        =       TFCvtForImageClassification(A     )
   UpperCAmelCase__       :	Union[str, Any]        =       model(A       ,labels=A       ,training=A     )
   self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_labels)     )




  def  __lowercase							(	self		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Union[str, Any]        =       self.prepare_config_and_inputs()
   UpperCAmelCase__ ,   UpperCAmelCase__ ,   UpperCAmelCase__       :	List[str]        =       config_and_inputs
   UpperCAmelCase__       :	List[str]        =       {"""pixel_values""": pixel_values}
   return config, inputs_dict







@require_tf
class   __lowercase      (     __lowerCamelCase						,      __lowerCamelCase						,      unittest.TestCase       ):
  snake_case_												=							(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
  snake_case_												=							(
      {"""feature-extraction""": TFCvtModel, """image-classification""": TFCvtForImageClassification}
      if is_tf_available()
      else {}
  )
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False
  snake_case_												=							False

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       TFCvtModelTester(self     )
   UpperCAmelCase__       :	Union[str, Any]        =       TFCvtConfigTester(self       ,config_class=A       ,has_text_modality=A       ,hidden_size=37     )

  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   self.config_tester.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()

  @unittest.skip(reason="""Cvt does not output attentions"""     )
  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''
   pass

  @unittest.skip(reason="""Cvt does not use inputs_embeds"""     )
  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   pass

  @unittest.skip(reason="""Cvt does not support input and output embeddings"""     )
  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   pass

  @unittest.skipIf(
      not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""     )     ) == 0       ,reason="""TF does not support backprop for grouped convolutions on CPU."""       ,)
  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   super().test_dataset_conversion()

  @unittest.skipIf(
      not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""     )     ) == 0       ,reason="""TF does not support backprop for grouped convolutions on CPU."""       ,)
  @slow
  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   super().test_keras_fit()

  @unittest.skip(reason="""Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8"""     )
  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       tf.keras.mixed_precision.Policy("""mixed_float16"""     )
   tf.keras.mixed_precision.set_global_policy(A     )
   super().test_keras_fit()
   tf.keras.mixed_precision.set_global_policy("""float32"""     )

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__ ,   UpperCAmelCase__       :	List[Any]        =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	int        =       model_class(A     )
    UpperCAmelCase__       :	List[Any]        =       inspect.signature(model.call     )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    UpperCAmelCase__       :	Optional[int]        =       [*signature.parameters.keys()]

    UpperCAmelCase__       :	Dict        =       ["""pixel_values"""]
    self.assertListEqual(arg_names[:1]       ,A     )

  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   def check_hidden_states_output(A		:      List[Any]       ,A		:      Tuple       ,A		:      int     ):
    UpperCAmelCase__       :	Any        =       model_class(A     )

    UpperCAmelCase__       :	Any        =       model(**self._prepare_for_class(A       ,A     )     )
    UpperCAmelCase__       :	List[str]        =       outputs.hidden_states

    UpperCAmelCase__       :	int        =       len(self.model_tester.depth     )
    self.assertEqual(len(A     )       ,A     )

    # verify the first hidden states (first block)
    self.assertListEqual(
        list(hidden_states[0].shape[-3:]     )       ,[
            self.model_tester.embed_dim[0],
            self.model_tester.image_size // 4,
            self.model_tester.image_size // 4,
        ]       ,)

   UpperCAmelCase__ ,   UpperCAmelCase__       :	str        =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__       :	Union[str, Any]        =       True
    check_hidden_states_output(A       ,A       ,A     )

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    UpperCAmelCase__       :	List[str]        =       True

    check_hidden_states_output(A       ,A       ,A     )

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*A     )

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*A     )




  @slow
  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    UpperCAmelCase__       :	List[Any]        =       TFCvtModel.from_pretrained(A     )
    self.assertIsNotNone(A     )







def 	lowerCAmelCase    (       ):
 '''simple docstring'''





 UpperCAmelCase__       :	Optional[Any]        =       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
 return image







@require_tf
@require_vision
class   __lowercase      (     unittest.TestCase       ):

  @cached_property
  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )




  @slow
  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )

   UpperCAmelCase__       :	Tuple        =       self.default_image_processor
   UpperCAmelCase__       :	Optional[Any]        =       prepare_img()
   UpperCAmelCase__       :	Optional[int]        =       image_processor(images=A       ,return_tensors="""tf"""     )

   # forward pass
   UpperCAmelCase__       :	Union[str, Any]        =       model(**A     )

   # verify the logits
   UpperCAmelCase__       :	List[str]        =       tf.TensorShape((1, 1_000)     )
   self.assertEqual(outputs.logits.shape       ,A     )

   UpperCAmelCase__       :	Any        =       tf.constant([0.9_2_8_5, 0.9_0_1_5, -0.3_1_5_0]     )
   self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()       ,A       ,atol=1e-4     )     )











'''simple docstring'''


import argparse
import json
import os

import tensorstore as ts
import torch
from flax import serialization
from flax.traverse_util import flatten_dict, unflatten_dict
from tensorflow.io import gfile

from transformers.modeling_utils import dtype_byte_size
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
    rename_keys,
)
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
from transformers.utils.hub import convert_file_size_to_int



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					List[str] ,    UpperCamelCase__    :					List[str]				):




     """simple docstring"""


     if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
          # expert layer
          a_     :							Optional[int]     						=  flax_key_tuple[:-1] + ("""weight""",)
          a_     :							Optional[int]     						=  torch.permute(UpperCamelCase__ ,    (0, 2, 1)				)
     elif flax_key_tuple[-1] == "kernel" and ".".join(UpperCamelCase__				):
          # linear layer
          a_     :							Any     						=  flax_key_tuple[:-1] + ("""weight""",)
          a_     :							Optional[int]     						=  flax_tensor.T
     elif flax_key_tuple[-1] in ["scale", "embedding"]:
          a_     :							Any     						=  flax_key_tuple[:-1] + ("""weight""",)

     return flax_key_tuple, flax_tensor



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					Union[str, Any] ,    UpperCamelCase__    :					str ,    UpperCamelCase__    :					Any				):




     """simple docstring"""


     if "metadata" in layer:
          a_     :							Optional[Any]     						=  layer.split("""metadata"""				)
          a_     :							List[str]     						=  """""".join(split_layer[0]				)[:-1]
          a_     :							List[str]     						=  [tuple(("""metadata""" + split_layer[1]).split("""/"""				)				)]
     elif "kvstore" in layer:
          a_     :							Union[str, Any]     						=  layer.split("""kvstore"""				)
          a_     :							Tuple     						=  """""".join(split_layer[0]				)[:-1]
          a_     :							Union[str, Any]     						=  [tuple(("""kvstore""" + split_layer[1]).split("""/"""				)				)]

     else:
          a_     :							List[str]     						=  layer.split("""/"""				)
          a_     :							Union[str, Any]     						=  """/""".join(split_layer[:-1]				)
          a_     :							Union[str, Any]     						=  (split_layer[-1],)

     if "kvstore/path" in layer:
          a_     :							Any     						=  F"{switch_checkpoint_path}/{checkpoint_info[layer]}"
     elif "kvstore/driver" in layer:
          a_     :							Optional[int]     						=  """file"""
     else:
          a_     :							str     						=  checkpoint_info[layer]

     return curr_real_layer_name, split_layer, content



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					Union[str, Any] ,    UpperCamelCase__    :					Any				):




     """simple docstring"""


     a_     :							int     						=  rename_keys(UpperCamelCase__				)
     a_     :							Any     						=  {}
     for k, v in current_block.items():
          a_     :							Union[str, Any]     						=  v
     a_     :							List[str]     						=  new_current_block
     torch.save(UpperCamelCase__ ,    UpperCamelCase__				)



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					List[Any] ,    UpperCamelCase__    :					Any ,    UpperCamelCase__    :					Optional[int] ,    UpperCamelCase__    :					List[str] ,    UpperCamelCase__    :					str = WEIGHTS_NAME				):




     """simple docstring"""


     a_     :							str     						=  convert_file_size_to_int(UpperCamelCase__				)
     a_     :							str     						=  []
     a_     :							Optional[int]     						=  {}
     a_     :							Optional[int]     						=  0
     a_     :							Any     						=  0

     os.makedirs(UpperCamelCase__ ,    exist_ok=UpperCamelCase__				)
     with gfile.GFile(switch_checkpoint_path + """/checkpoint""" ,    """rb"""				) as fp:
          a_     :							str     						=  serialization.msgpack_restore(fp.read()				)["""optimizer"""]["""target"""]
          a_     :							List[str]     						=  flatten_dict(UpperCamelCase__ ,    sep="""/"""				)

     a_     :							Dict     						=  {}
     for layer in checkpoint_info.keys():
          a_						,     a_						,     a_     :							List[Any]     						=  get_key_and_tensorstore_dict(
              UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__				)
          if curr_real_layer_name in all_layers:
               a_     :							int     						=  content
          else:
               a_     :							Tuple     						=  {split_layer[-1]: content}

     for key in all_layers.keys():
          # open tensorstore file
          a_     :							List[str]     						=  ts.open(unflatten_dict(all_layers[key]				)				).result().read().result()
          a_     :							List[Any]     						=  torch.tensor(UpperCamelCase__				)
          a_     :							Dict     						=  raw_weights.numel() * dtype_byte_size(raw_weights.dtype				)

          # use the renaming pattern from the small conversion scripts
          a_						,     a_     :							Tuple     						=  rename_base_flax_keys(tuple(key.split("""/"""				)				) ,    UpperCamelCase__				)
          a_     :							Dict     						=  """/""".join(UpperCamelCase__				)

          # If this weight is going to tip up over the maximal size, we split.
          if current_block_size + weight_size > max_shard_size:
               a_     :							Any     						=  os.path.join(
                   UpperCamelCase__ ,    weights_name.replace(""".bin""" ,    F"-{len(UpperCamelCase__				)+1:05d}-of-???.bin"				)				)
               rename_and_save_block(UpperCamelCase__ ,    UpperCamelCase__				)
               sharded_state_dicts.append(current_block.keys()				)
               del current_block
               a_     :							Optional[int]     						=  {}
               a_     :							Union[str, Any]     						=  0

          a_     :							List[str]     						=  raw_weights.to(getattr(UpperCamelCase__ ,    UpperCamelCase__				)				)
          current_block_size += weight_size
          total_size += weight_size

     # Add the last block
     a_     :							Any     						=  os.path.join(UpperCamelCase__ ,    weights_name.replace(""".bin""" ,    F"-{len(UpperCamelCase__				)+1:05d}-of-???.bin"				)				)
     rename_and_save_block(UpperCamelCase__ ,    UpperCamelCase__				)
     sharded_state_dicts.append(current_block.keys()				)

     # If we only have one shard, we return it
     if len(UpperCamelCase__				) == 1:
          return {weights_name: sharded_state_dicts[0]}, None

     # Otherwise, let's build the index
     a_     :							Tuple     						=  {}
     a_     :							List[Any]     						=  {}
     for idx, shard in enumerate(UpperCamelCase__				):
          a_     :							str     						=  weights_name.replace(
              """.bin""" ,    F"-{idx+1:05d}-of-{len(UpperCamelCase__				):05d}.bin"				)  # len(sharded_state_dicts):05d}
          a_     :							Union[str, Any]     						=  os.path.join(UpperCamelCase__ ,    weights_name.replace(""".bin""" ,    F"-{idx+1:05d}-of-???.bin"				)				)
          os.rename(UpperCamelCase__ ,    os.path.join(UpperCamelCase__ ,    UpperCamelCase__				)				)
          a_     :							str     						=  shard
          for key in shard:
               a_     :							Union[str, Any]     						=  shard_file

    # Add the metadata
     a_     :							List[str]     						=  {"""total_size""": total_size}
     a_     :							str     						=  {"""metadata""": metadata, """weight_map""": weight_map}

     with open(os.path.join(UpperCamelCase__ ,    UpperCamelCase__				) ,    """w""" ,    encoding="""utf-8"""				) as f:
          a_     :							Any     						=  json.dumps(UpperCamelCase__ ,    indent=2 ,    sort_keys=UpperCamelCase__				) + """\n"""
          f.write(UpperCamelCase__				)

     return metadata, index


if __name__ == "__main__":
 lowerCAmelCase_      :  List[Any]				     =  argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--switch_t5x_checkpoint_path',
     default='/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600',
     type=str,
     required=False,
     help='Path to a directory containing a folder per layer. Follows the original Google format.',
 )
 parser.add_argument('--max_shard_size', default='10GB', required=False, help='Max shard size')
 parser.add_argument('--dtype', default='bfloat16', type=str, required=False, help='dtype of the saved model')
 parser.add_argument(
     '--pytorch_dump_folder_path',
     default='/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted',
     type=str,
     required=False,
     help='Path to the output pytorch model.',
 )
 lowerCAmelCase_      :  List[Any]				     =  parser.parse_args()
 shard_on_the_fly(
     args.switch_tax_checkpoint_path,
     args.pytorch_dump_folder_path,
     args.max_shard_size,
     args.dtype,
 )



def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

     a_     :							Union[str, Any]     						=  SwitchTransformersConfig.from_pretrained("""google/switch-base-8"""				)
     config.save_pretrained("""/home/arthur_huggingface_co/transformers/switch_converted"""				)
     a_     :							List[Any]     						=  SwitchTransformersForConditionalGeneration.from_pretrained(
         """/home/arthur_huggingface_co/transformers/switch_converted""" ,    device_map="""auto"""				)

     a_     :							Union[str, Any]     						=  TaTokenizer.from_pretrained("""t5-small"""				)
     a_     :							Optional[Any]     						=  """A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>."""

     a_     :							List[Any]     						=  tokenizer(UpperCamelCase__ ,    return_tensors="""pt"""				).input_ids
     a_     :							str     						=  model.generate(UpperCamelCase__ ,    decoder_start_token_id=0				)
     print(tokenizer.decode(out[0]				)				)
