

'''simple docstring'''






import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	Dict    )							->   int:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_      :			Dict     =       ConvNextVaConfig()

      if "atto" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Any     =       [2, 2, 6, 2]
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       [4_0, 8_0, 1_6_0, 3_2_0]
      if "femto" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       [2, 2, 6, 2]
            SCREAMING_SNAKE_CASE_      :			Dict     =       [4_8, 9_6, 1_9_2, 3_8_4]
      if "pico" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Dict     =       [2, 2, 6, 2]
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       [6_4, 1_2_8, 2_5_6, 5_1_2]
      if "nano" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       [2, 2, 8, 2]
            SCREAMING_SNAKE_CASE_      :			List[str]     =       [8_0, 1_6_0, 3_2_0, 6_4_0]
      if "tiny" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       [3, 3, 9, 3]
            SCREAMING_SNAKE_CASE_      :			List[str]     =       [9_6, 1_9_2, 3_8_4, 7_6_8]
      if "base" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       [3, 3, 2_7, 3]
            SCREAMING_SNAKE_CASE_      :			str     =       [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
      if "large" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       [3, 3, 2_7, 3]
            SCREAMING_SNAKE_CASE_      :			Dict     =       [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
      if "huge" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Dict     =       [3, 3, 2_7, 3]
            SCREAMING_SNAKE_CASE_      :			Dict     =       [3_5_2, 7_0_4, 1_4_0_8, 2_8_1_6]

      SCREAMING_SNAKE_CASE_      :			str     =       1_0_0_0
      SCREAMING_SNAKE_CASE_      :			str     =       "imagenet-1k-id2label.json"
      SCREAMING_SNAKE_CASE_      :			List[str]     =       (1, 1_0_0_0)

      SCREAMING_SNAKE_CASE_      :			List[Any]     =       "huggingface/label-files"
      SCREAMING_SNAKE_CASE_      :			Dict     =       num_labels
      SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_  ,  SCREAMING_SNAKE_CASE_  ,  repo_type="dataset"    )  ,  "r"    )    )
      SCREAMING_SNAKE_CASE_      :			int     =       {int(SCREAMING_SNAKE_CASE_    ): v for k, v in idalabel.items()}

      SCREAMING_SNAKE_CASE_      :			List[str]     =       idalabel
      SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       {v: k for k, v in idalabel.items()}
      SCREAMING_SNAKE_CASE_      :			List[str]     =       hidden_sizes
      SCREAMING_SNAKE_CASE_      :			Dict     =       depths

      return config, expected_shape


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	Dict    )							->   str:
      """simple docstring"""

      if "downsample_layers.0.0" in name:
            SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       name.replace("downsample_layers.0.0"  ,  "embeddings.patch_embeddings"    )
      if "downsample_layers.0.1" in name:
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       name.replace("downsample_layers.0.1"  ,  "embeddings.norm"    )  # we rename to layernorm later on
      if "downsample_layers.1.0" in name:
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       name.replace("downsample_layers.1.0"  ,  "stages.1.downsampling_layer.0"    )
      if "downsample_layers.1.1" in name:
            SCREAMING_SNAKE_CASE_      :			Dict     =       name.replace("downsample_layers.1.1"  ,  "stages.1.downsampling_layer.1"    )
      if "downsample_layers.2.0" in name:
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       name.replace("downsample_layers.2.0"  ,  "stages.2.downsampling_layer.0"    )
      if "downsample_layers.2.1" in name:
            SCREAMING_SNAKE_CASE_      :			int     =       name.replace("downsample_layers.2.1"  ,  "stages.2.downsampling_layer.1"    )
      if "downsample_layers.3.0" in name:
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       name.replace("downsample_layers.3.0"  ,  "stages.3.downsampling_layer.0"    )
      if "downsample_layers.3.1" in name:
            SCREAMING_SNAKE_CASE_      :			Tuple     =       name.replace("downsample_layers.3.1"  ,  "stages.3.downsampling_layer.1"    )
      if "stages" in name and "downsampling_layer" not in name:
            # stages.0.0. for instance should be renamed to stages.0.layers.0.
            SCREAMING_SNAKE_CASE_      :			str     =       name[: len("stages.0"    )] + ".layers" + name[len("stages.0"    ) :]
      if "gamma" in name:
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       name.replace("gamma"  ,  "weight"    )
      if "beta" in name:
            SCREAMING_SNAKE_CASE_      :			Dict     =       name.replace("beta"  ,  "bias"    )
      if "stages" in name:
            SCREAMING_SNAKE_CASE_      :			Tuple     =       name.replace("stages"  ,  "encoder.stages"    )
      if "norm" in name:
            SCREAMING_SNAKE_CASE_      :			Optional[int]     =       name.replace("norm"  ,  "layernorm"    )
      if "head" in name:
            SCREAMING_SNAKE_CASE_      :			Any     =       name.replace("head"  ,  "classifier"    )

      return name


def 					__lowerCamelCase			(      )							->   str:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_      :			List[Any]     =       "http://images.cocodataset.org/val2017/000000039769.jpg"
      SCREAMING_SNAKE_CASE_      :			str     =       Image.open(requests.get(SCREAMING_SNAKE_CASE_  ,  stream=SCREAMING_SNAKE_CASE_    ).raw    )
      return im


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	Any    )							->   Any:
      """simple docstring"""

      if "224" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Any     =       2_2_4
            SCREAMING_SNAKE_CASE_      :			int     =       2_2_4 / 2_5_6
      elif "384" in checkpoint_url:
            SCREAMING_SNAKE_CASE_      :			Dict     =       3_8_4
            SCREAMING_SNAKE_CASE_      :			str     =       None
      else:
            SCREAMING_SNAKE_CASE_      :			Dict     =       5_1_2
            SCREAMING_SNAKE_CASE_      :			List[str]     =       None

      return ConvNextImageProcessor(
          size=SCREAMING_SNAKE_CASE_  ,  crop_pct=SCREAMING_SNAKE_CASE_  ,  image_mean=[0.485, 0.456, 0.406]  ,  image_std=[0.229, 0.224, 0.225]  ,  resample=PILImageResampling.BICUBIC  ,  )


@torch.no_grad()
def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	Tuple  ,  SCREAMING_SNAKE_CASE_     :	Any  ,  SCREAMING_SNAKE_CASE_     :	Optional[Any]  ,  SCREAMING_SNAKE_CASE_     :	Union[str, Any]    )							->   str:
      """simple docstring"""

      print("Downloading original model from checkpoint..."    )
      # define ConvNeXTV2 configuration based on URL
      SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Tuple     =       get_convnextva_config(SCREAMING_SNAKE_CASE_    )
      # load original state_dict from URL
      SCREAMING_SNAKE_CASE_      :			Tuple     =       torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_    )["model"]

      print("Converting model parameters..."    )
      # rename keys
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE_      :			Any     =       state_dict.pop(SCREAMING_SNAKE_CASE_    )
            SCREAMING_SNAKE_CASE_      :			List[str]     =       val
      # add prefix to all keys expect classifier head
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE_      :			Any     =       state_dict.pop(SCREAMING_SNAKE_CASE_    )
            if not key.startswith("classifier"    ):
                  SCREAMING_SNAKE_CASE_      :			int     =       "convnextv2." + key
            SCREAMING_SNAKE_CASE_      :			Tuple     =       val

      # load HuggingFace model
      SCREAMING_SNAKE_CASE_      :			List[Any]     =       ConvNextVaForImageClassification(SCREAMING_SNAKE_CASE_    )
      model.load_state_dict(SCREAMING_SNAKE_CASE_    )
      model.eval()

      # Check outputs on an image, prepared by ConvNextImageProcessor
      SCREAMING_SNAKE_CASE_      :			Optional[int]     =       convert_preprocessor(SCREAMING_SNAKE_CASE_    )
      SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       preprocessor(images=prepare_img()  ,  return_tensors="pt"    )
      SCREAMING_SNAKE_CASE_      :			str     =       model(**SCREAMING_SNAKE_CASE_    ).logits

      # note: the logits below were obtained without center cropping
      if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Tuple     =       torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			str     =       torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			str     =       torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Any     =       torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			str     =       torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Any     =       torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE_      :			str     =       torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Dict     =       torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]    )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
            SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]    )
      else:
            raise ValueError(F"Unknown URL: {checkpoint_url}"    )

      assert torch.allclose(logits[0, :5]  ,  SCREAMING_SNAKE_CASE_  ,  atol=1E-3    )
      assert logits.shape == expected_shape
      print("Model outputs match the original results!"    )

      if save_model:
            print("Saving model to local..."    )
            # Create folder to save model
            if not os.path.isdir(SCREAMING_SNAKE_CASE_    ):
                  os.mkdir(SCREAMING_SNAKE_CASE_    )

            model.save_pretrained(SCREAMING_SNAKE_CASE_    )
            preprocessor.save_pretrained(SCREAMING_SNAKE_CASE_    )

      SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       "convnextv2"
      if "atto" in checkpoint_url:
            model_name += "-atto"
      if "femto" in checkpoint_url:
            model_name += "-femto"
      if "pico" in checkpoint_url:
            model_name += "-pico"
      if "nano" in checkpoint_url:
            model_name += "-nano"
      elif "tiny" in checkpoint_url:
            model_name += "-tiny"
      elif "base" in checkpoint_url:
            model_name += "-base"
      elif "large" in checkpoint_url:
            model_name += "-large"
      elif "huge" in checkpoint_url:
            model_name += "-huge"
      if "22k" in checkpoint_url and "1k" not in checkpoint_url:
            model_name += "-22k"
      elif "22k" in checkpoint_url and "1k" in checkpoint_url:
            model_name += "-22k-1k"
      elif "1k" in checkpoint_url:
            model_name += "-1k"
      if "224" in checkpoint_url:
            model_name += "-224"
      elif "384" in checkpoint_url:
            model_name += "-384"
      elif "512" in checkpoint_url:
            model_name += "-512"

      if push_to_hub:
            print(F"Pushing {model_name} to the hub..."    )
            model.push_to_hub(SCREAMING_SNAKE_CASE_    )
            preprocessor.push_to_hub(SCREAMING_SNAKE_CASE_    )


if __name__ == "__main__":
      snake_case_       =       argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '--checkpoint_url',
          default='https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt',
          type=str,
          help='URL of the original ConvNeXTV2 checkpoint you\'d like to convert.',
      )
      parser.add_argument(
          '--pytorch_dump_folder_path',
          default='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 preprocessor to the hub')

      snake_case_       =       parser.parse_args()
      convert_convnextva_checkpoint(
          args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
      )



"""simple docstring"""




import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    CpmAntConfig,
	    CpmAntForCausalLM,
	    CpmAntModel,
	    CpmAntTokenizer,
	)




@require_torch
class      UpperCamelCase	:

				"""simple docstring"""



				def __init__(				self		,	lowercase__		,	lowercase__=2		,	lowercase__=8		,	lowercase__=True		,	lowercase__=False		,	lowercase__=False		,	lowercase__=False		,	lowercase__=False		,	lowercase__=99		,	lowercase__=32		,	lowercase__=3		,	lowercase__=4		,	lowercase__=37		,	lowercase__=32		,	lowercase__=128		,	lowercase__=8		,	lowercase__=8		,	lowercase__=8		,	lowercase__=1.0		,	lowercase__=True		,	)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	parent
									SCREAMING_SNAKE_CASE				=	batch_size
									SCREAMING_SNAKE_CASE				=	seq_length
									SCREAMING_SNAKE_CASE				=	is_training
									SCREAMING_SNAKE_CASE				=	use_token_type_ids
									SCREAMING_SNAKE_CASE				=	use_input_mask
									SCREAMING_SNAKE_CASE				=	use_labels
									SCREAMING_SNAKE_CASE				=	use_mc_token_ids
									SCREAMING_SNAKE_CASE				=	vocab_size
									SCREAMING_SNAKE_CASE				=	hidden_size
									SCREAMING_SNAKE_CASE				=	num_hidden_layers
									SCREAMING_SNAKE_CASE				=	num_attention_heads
									SCREAMING_SNAKE_CASE				=	intermediate_size
									SCREAMING_SNAKE_CASE				=	num_buckets
									SCREAMING_SNAKE_CASE				=	max_distance
									SCREAMING_SNAKE_CASE				=	prompt_length
									SCREAMING_SNAKE_CASE				=	prompt_types
									SCREAMING_SNAKE_CASE				=	segment_types
									SCREAMING_SNAKE_CASE				=	init_std
									SCREAMING_SNAKE_CASE				=	return_dict


				def 							A					(				self			)    ->    Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	{}
									SCREAMING_SNAKE_CASE				=	ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size			).type(torch.intaa			)
									SCREAMING_SNAKE_CASE				=	False

									SCREAMING_SNAKE_CASE				=	self.get_config()

									return (config, input_ids)


				def 							A					(				self			)    ->    List[Any]:
									"""simple docstring"""







									return CpmAntConfig(
									    vocab_size=self.vocab_size		,	hidden_size=self.hidden_size		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	dim_ff=self.intermediate_size		,	position_bias_num_buckets=self.num_buckets		,	position_bias_max_distance=self.max_distance		,	prompt_types=self.prompt_types		,	prompt_length=self.prompt_length		,	segment_types=self.segment_types		,	use_cache=lowercase__		,	init_std=self.init_std		,	return_dict=self.return_dict		,	)


				def 							A					(				self		,	lowercase__		,	lowercase__		,	*lowercase__			)    ->    Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntModel(config=lowercase__			)
									model.to(lowercase__			)
									model.eval()

									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).last_hidden_state

									self.parent.assertEqual(hidden_states.shape		,	(self.batch_size, self.seq_length, config.hidden_size)			)


				def 							A					(				self		,	lowercase__		,	lowercase__		,	*lowercase__			)    ->    int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM(lowercase__			)
									model.to(lowercase__			)
									SCREAMING_SNAKE_CASE				=	input_ids['input_ids'].to(lowercase__			)
									model.eval()

									SCREAMING_SNAKE_CASE				=	model(**lowercase__			)
									self.parent.assertEqual(
									    model_output.logits.shape		,	(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)		,	)




				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.prepare_config_and_inputs()
									return config, inputs_dict






@require_torch
class      UpperCamelCase	(       __magic_name__  ,      __magic_name__  ,      unittest.TestCase      ):

				"""simple docstring"""



				UpperCAmelCase_				:      List[str]						  =				(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
				UpperCAmelCase_				:      List[Any]						  =				(
				    {"feature-extraction": CpmAntModel, "text-generation": CpmAntForCausalLM} if is_torch_available() else {}
				)

				UpperCAmelCase_				:      List[Any]						  =				False
				UpperCAmelCase_				:      Optional[int]						  =				False
				UpperCAmelCase_				:      Optional[Any]						  =				False
				UpperCAmelCase_				:      List[str]						  =				False
				UpperCAmelCase_				:      Dict						  =				False


				def 							A					(				self			)    ->    Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntModelTester(self			)
									SCREAMING_SNAKE_CASE				=	ConfigTester(self		,	config_class=lowercase__			)


				def 							A					(				self			)    ->    Tuple:
									"""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.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()


				def 							A					(				self			)    ->    List[str]:
									"""simple docstring"""







									unittest.skip('CPMAnt doesn\'t support input_embeds.'			)(self.test_inputs_embeds			)


				def 							A					(				self			)    ->    Optional[int]:
									"""simple docstring"""







									unittest.skip(
									    'CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'			)(self.test_retain_grad_hidden_states_attentions			)


				def 							A					(				self			)    ->    List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_cpmant_model(lowercase__		,	lowercase__			)




				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_lm_head_model(lowercase__		,	lowercase__			)






@require_torch
class      UpperCamelCase	(       unittest.TestCase      ):

				"""simple docstring"""



				@tooslow
				def 							A					(				self			)    ->    Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'今天天气真好！'
									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntModel.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).last_hidden_state

									SCREAMING_SNAKE_CASE				=	torch.tensor(
									    [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]		,	)
									self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,	lowercase__		,	atol=1E-2			)			)






@require_torch
class      UpperCamelCase	(       unittest.TestCase      ):

				"""simple docstring"""



				@tooslow
				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'今天天气真好！'
									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).logits

									SCREAMING_SNAKE_CASE				=	torch.tensor(
									    [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]		,	)
									self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,	lowercase__		,	atol=1E-2			)			)


				@tooslow
				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	'今天天气不错，'
									SCREAMING_SNAKE_CASE				=	'今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model.generate(**lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer.batch_decode(lowercase__			)
									self.assertEqual(lowercase__		,	lowercase__			)




				@tooslow
				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	['今天天气不错，', '新年快乐，万事如意！']
									SCREAMING_SNAKE_CASE				=	[
									    '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
									    '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
									]
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'		,	padding=lowercase__			)
									SCREAMING_SNAKE_CASE				=	model.generate(**lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer.batch_decode(lowercase__			)
									self.assertEqual(lowercase__		,	lowercase__			)
