



'''simple docstring'''



import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
	import jax
	import jax.numpy as jnp

	from transformers.modeling_flax_pytorch_utils import (
	    convert_pytorch_state_dict_to_flax,
	    load_flax_weights_in_pytorch_model,
	)
	from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
	import torch






class 					UpperCAmelCase__	:



					def __init__(    self			,				UpperCamelCase			,				UpperCamelCase=14			,				UpperCamelCase=7			,				UpperCamelCase=True			,				UpperCamelCase=True			,				UpperCamelCase=False			,				UpperCamelCase=True			,				UpperCamelCase=99			,				UpperCamelCase=32			,				UpperCamelCase=4			,				UpperCamelCase=4			,				UpperCamelCase=4			,				UpperCamelCase=37			,				UpperCamelCase="gelu"			,				UpperCamelCase=0.1			,				UpperCamelCase=0.1			,				UpperCamelCase=512			,				UpperCamelCase=0.02			,				)     ->			List[Any]:
						__lowerCAmelCase												=  parent
						__lowerCAmelCase												=  batch_size
						__lowerCAmelCase												=  seq_length
						__lowerCAmelCase												=  is_training
						__lowerCAmelCase												=  use_input_mask
						__lowerCAmelCase												=  use_token_type_ids
						__lowerCAmelCase												=  use_labels
						__lowerCAmelCase												=  vocab_size
						__lowerCAmelCase												=  hidden_size
						__lowerCAmelCase												=  rotary_dim
						__lowerCAmelCase												=  num_hidden_layers
						__lowerCAmelCase												=  num_attention_heads
						__lowerCAmelCase												=  intermediate_size
						__lowerCAmelCase												=  hidden_act
						__lowerCAmelCase												=  hidden_dropout_prob
						__lowerCAmelCase												=  attention_probs_dropout_prob
						__lowerCAmelCase												=  max_position_embeddings
						__lowerCAmelCase												=  initializer_range
						__lowerCAmelCase												=  None
						__lowerCAmelCase												=  vocab_size - 1
						__lowerCAmelCase												=  vocab_size - 1
						__lowerCAmelCase												=  vocab_size - 1



					def        UpperCAmelCase_ (    self     )     ->			Tuple:
						__lowerCAmelCase												=  ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size     )

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

						__lowerCAmelCase												=  GPTJConfig(
						    vocab_size=self.vocab_size			,				n_embd=self.hidden_size			,				n_layer=self.num_hidden_layers			,				n_head=self.num_attention_heads			,				n_positions=self.max_position_embeddings			,				use_cache=UpperCamelCase			,				bos_token_id=self.bos_token_id			,				eos_token_id=self.eos_token_id			,				pad_token_id=self.pad_token_id			,				rotary_dim=self.rotary_dim			,				)

						return (config, input_ids, input_mask)



					def        UpperCAmelCase_ (    self     )     ->			Optional[Any]:
						__lowerCAmelCase												=  self.prepare_config_and_inputs()
						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  config_and_inputs
						__lowerCAmelCase												=  {"input_ids": input_ids, "attention_mask": attention_mask}
						return config, inputs_dict



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Optional[int]:
						__lowerCAmelCase												=  20
						__lowerCAmelCase												=  model_class_name(UpperCamelCase     )

						__lowerCAmelCase												=  model.init_cache(input_ids.shape[0]			,				UpperCamelCase     )
						__lowerCAmelCase												=  jnp.ones((input_ids.shape[0], max_decoder_length)			,				dtype="i4"     )

						__lowerCAmelCase												=  jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1     )[None, :]			,				(input_ids.shape[0], input_ids.shape[-1] - 1)     )
						__lowerCAmelCase												=  model(
						    input_ids[:, :-1]			,				attention_mask=UpperCamelCase			,				past_key_values=UpperCamelCase			,				position_ids=UpperCamelCase			,				)

						__lowerCAmelCase												=  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,				dtype="i4"     )
						__lowerCAmelCase												=  model(
						    input_ids[:, -1:]			,				attention_mask=UpperCamelCase			,				past_key_values=outputs_cache.past_key_values			,				position_ids=UpperCamelCase			,				)

						__lowerCAmelCase												=  model(UpperCamelCase     )

						__lowerCAmelCase												=  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
						self.parent.assertTrue(diff < 1E-3			,				msg=F'''Max diff is {diff}'''     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Tuple:
						__lowerCAmelCase												=  20
						__lowerCAmelCase												=  model_class_name(UpperCamelCase     )

						__lowerCAmelCase												=  jnp.concatenate(
						    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])     )]			,				axis=-1			,				)

						__lowerCAmelCase												=  model.init_cache(input_ids.shape[0]			,				UpperCamelCase     )
						__lowerCAmelCase												=  jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1     )[None, :]			,				(input_ids.shape[0], input_ids.shape[-1] - 1)     )

						__lowerCAmelCase												=  model(
						    input_ids[:, :-1]			,				attention_mask=UpperCamelCase			,				past_key_values=UpperCamelCase			,				position_ids=UpperCamelCase			,				)
						__lowerCAmelCase												=  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]			,				dtype="i4"     )
						__lowerCAmelCase												=  model(
						    input_ids[:, -1:]			,				past_key_values=outputs_cache.past_key_values			,				attention_mask=UpperCamelCase			,				position_ids=UpperCamelCase			,				)

						__lowerCAmelCase												=  model(UpperCamelCase			,				attention_mask=UpperCamelCase     )

						__lowerCAmelCase												=  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
						self.parent.assertTrue(diff < 1E-3			,				msg=F'''Max diff is {diff}'''     )






@require_flax
class 					UpperCAmelCase__	(		UpperCamelCase__							,       UpperCamelCase__							,       unittest.TestCase ):
					a      :     Union[str, Any]        =							(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
					a      :     Any        =							(FlaxGPTJForCausalLM,) if is_flax_available() else ()



					def        UpperCAmelCase_ (    self     )     ->			int:
						__lowerCAmelCase												=  FlaxGPTJModelTester(self     )



					def        UpperCAmelCase_ (    self     )     ->			str:
						for model_class_name in self.all_model_classes:
							__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.check_use_cache_forward(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )



					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						for model_class_name in self.all_model_classes:
							__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.check_use_cache_forward_with_attn_mask(
							    UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )



					@tooslow
					def        UpperCAmelCase_ (    self     )     ->			str:
						__lowerCAmelCase												=  GPTaTokenizer.from_pretrained("gpt2"			,				pad_token="<|endoftext|>"			,				padding_side="left"     )
						__lowerCAmelCase												=  tokenizer(["Hello this is a long string", "Hey"]			,				return_tensors="np"			,				padding=UpperCamelCase			,				truncation=UpperCamelCase     )

						__lowerCAmelCase												=  FlaxGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"     )
						__lowerCAmelCase												=  False
						__lowerCAmelCase												=  model.config.eos_token_id

						__lowerCAmelCase												=  jax.jit(model.generate     )

						__lowerCAmelCase												=  jit_generate(
						    inputs["input_ids"]			,				attention_mask=inputs["attention_mask"]			,				pad_token_id=tokenizer.pad_token_id     ).sequences

						__lowerCAmelCase												=  tokenizer.batch_decode(UpperCamelCase			,				skip_special_tokens=UpperCamelCase     )

						__lowerCAmelCase												=  [
						    "Hello this is a long string of text.\n\nI'm trying to get the text of the",
						    "Hey, I'm a little late to the party. I'm going to",
						]

						self.assertListEqual(UpperCamelCase			,				UpperCamelCase     )



					@is_pt_flax_cross_test
					def        UpperCAmelCase_ (    self     )     ->			str:
						__lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__     ):
								# prepare inputs
								__lowerCAmelCase												=  self._prepare_for_class(UpperCamelCase			,				UpperCamelCase     )
								__lowerCAmelCase												=  {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

								# load corresponding PyTorch class
								__lowerCAmelCase												=  model_class.__name__[4:]  # Skip the "Flax" at the beginning
								__lowerCAmelCase												=  getattr(UpperCamelCase			,				UpperCamelCase     )

								__lowerCAmelCase							,       __lowerCAmelCase												=  pt_inputs["input_ids"].shape
								__lowerCAmelCase												=  np.random.randint(0			,				seq_length - 1			,				size=(batch_size,)     )
								for batch_idx, start_index in enumerate(UpperCamelCase     ):
									__lowerCAmelCase												=  0
									__lowerCAmelCase												=  1
									__lowerCAmelCase												=  0
									__lowerCAmelCase												=  1
								__lowerCAmelCase												=  pt_model_class(UpperCamelCase     ).eval()
								__lowerCAmelCase												=  model_class(UpperCamelCase			,				dtype=jnp.floataa     )

								__lowerCAmelCase												=  convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,				UpperCamelCase     )
								__lowerCAmelCase												=  fx_state

								with torch.no_grad():
									__lowerCAmelCase												=  pt_model(**UpperCamelCase     ).to_tuple()

								__lowerCAmelCase												=  fx_model(**UpperCamelCase     ).to_tuple()
								self.assertEqual(len(UpperCamelCase     )			,				len(UpperCamelCase     )			,				"Output lengths differ between Flax and PyTorch"     )
								for fx_output, pt_output in zip(UpperCamelCase			,				UpperCamelCase     ):
									self.assert_almost_equals(fx_output[:, -1]			,				pt_output[:, -1].numpy()			,				4E-2     )

								with tempfile.TemporaryDirectory() as tmpdirname:
									pt_model.save_pretrained(UpperCamelCase     )
									__lowerCAmelCase												=  model_class.from_pretrained(UpperCamelCase			,				from_pt=UpperCamelCase     )

								__lowerCAmelCase												=  fx_model_loaded(**UpperCamelCase     ).to_tuple()
								self.assertEqual(
								    len(UpperCamelCase     )			,				len(UpperCamelCase     )			,				"Output lengths differ between Flax and PyTorch"     )
								for fx_output_loaded, pt_output in zip(UpperCamelCase			,				UpperCamelCase     ):
									self.assert_almost_equals(fx_output_loaded[:, -1]			,				pt_output[:, -1].numpy()			,				4E-2     )



					@is_pt_flax_cross_test
					def        UpperCAmelCase_ (    self     )     ->			List[str]:
						__lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__     ):
								# prepare inputs
								__lowerCAmelCase												=  self._prepare_for_class(UpperCamelCase			,				UpperCamelCase     )
								__lowerCAmelCase												=  {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

								# load corresponding PyTorch class
								__lowerCAmelCase												=  model_class.__name__[4:]  # Skip the "Flax" at the beginning
								__lowerCAmelCase												=  getattr(UpperCamelCase			,				UpperCamelCase     )

								__lowerCAmelCase												=  pt_model_class(UpperCamelCase     ).eval()
								__lowerCAmelCase												=  model_class(UpperCamelCase			,				dtype=jnp.floataa     )

								__lowerCAmelCase												=  load_flax_weights_in_pytorch_model(UpperCamelCase			,				fx_model.params     )
								__lowerCAmelCase							,       __lowerCAmelCase												=  pt_inputs["input_ids"].shape
								__lowerCAmelCase												=  np.random.randint(0			,				seq_length - 1			,				size=(batch_size,)     )
								for batch_idx, start_index in enumerate(UpperCamelCase     ):
									__lowerCAmelCase												=  0
									__lowerCAmelCase												=  1
									__lowerCAmelCase												=  0
									__lowerCAmelCase												=  1

								# make sure weights are tied in PyTorch
								pt_model.tie_weights()

								with torch.no_grad():
									__lowerCAmelCase												=  pt_model(**UpperCamelCase     ).to_tuple()

								__lowerCAmelCase												=  fx_model(**UpperCamelCase     ).to_tuple()
								self.assertEqual(len(UpperCamelCase     )			,				len(UpperCamelCase     )			,				"Output lengths differ between Flax and PyTorch"     )
								for fx_output, pt_output in zip(UpperCamelCase			,				UpperCamelCase     ):
									self.assert_almost_equals(fx_output[:, -1]			,				pt_output[:, -1].numpy()			,				4E-2     )

								with tempfile.TemporaryDirectory() as tmpdirname:
									fx_model.save_pretrained(UpperCamelCase     )
									__lowerCAmelCase												=  pt_model_class.from_pretrained(UpperCamelCase			,				from_flax=UpperCamelCase     )

								with torch.no_grad():
									__lowerCAmelCase												=  pt_model_loaded(**UpperCamelCase     ).to_tuple()

								self.assertEqual(
								    len(UpperCamelCase     )			,				len(UpperCamelCase     )			,				"Output lengths differ between Flax and PyTorch"     )
								for fx_output, pt_output in zip(UpperCamelCase			,				UpperCamelCase     ):
									self.assert_almost_equals(fx_output[:, -1]			,				pt_output[:, -1].numpy()			,				4E-2     )



					@tooslow
					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						for model_class_name in self.all_model_classes:
							__lowerCAmelCase												=  model_class_name.from_pretrained("EleutherAI/gpt-j-6B"     )
							__lowerCAmelCase												=  model(np.ones((1, 1)     )     )
							self.assertIsNotNone(UpperCamelCase     )




'''simple docstring'''



import argparse
import json

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

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling



def   __lowerCAmelCase					(							lowerCamelCase   :						Any					):






	'''simple docstring'''




	__lowerCAmelCase												=  [2, 2, 6, 2] if "tiny" in model_name else [2, 2, 18, 2]
	__lowerCAmelCase												=  True if "large" in model_name or "huge" in model_name else False
	__lowerCAmelCase												=  True if "large" in model_name or "huge" in model_name else False
	__lowerCAmelCase												=  True if "large" in model_name or "huge" in model_name else False

	if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
		if "fl3" in model_name:
			__lowerCAmelCase												=  [3, 3, 3, 3]
			__lowerCAmelCase												=  [5, 5, 5, 5]
		elif "fl4" in model_name:
			__lowerCAmelCase												=  [4, 4, 4, 4]
			__lowerCAmelCase												=  [3, 3, 3, 3]

	if "tiny" in model_name or "small" in model_name or "base" in model_name:
		__lowerCAmelCase												=  [3, 3, 3, 3]
		if "lrf" in model_name:
			__lowerCAmelCase												=  [3, 3, 3, 3]
		else:
			__lowerCAmelCase												=  [2, 2, 2, 2]

	if "tiny" in model_name:
		__lowerCAmelCase												=  96
	elif "small" in model_name:
		__lowerCAmelCase												=  96
	elif "base" in model_name:
		__lowerCAmelCase												=  1_28
	elif "large" in model_name:
		__lowerCAmelCase												=  1_92
	elif "xlarge" in model_name:
		__lowerCAmelCase												=  2_56
	elif "huge" in model_name:
		__lowerCAmelCase												=  3_52

	# set label information
	__lowerCAmelCase												=  "huggingface/label-files"
	if "large" in model_name or "huge" in model_name:
		__lowerCAmelCase												=  "imagenet-22k-id2label.json"
	else:
		__lowerCAmelCase												=  "imagenet-1k-id2label.json"

	__lowerCAmelCase												=  json.load(open(hf_hub_download(lowerCamelCase				,		lowerCamelCase				,		repo_type="dataset"					)				,		"r"					)					)
	__lowerCAmelCase												=  {int(lowerCamelCase					): v for k, v in idalabel.items()}
	__lowerCAmelCase												=  {v: k for k, v in idalabel.items()}

	__lowerCAmelCase												=  FocalNetConfig(
	    embed_dim=lowerCamelCase				,		depths=lowerCamelCase				,		focal_levels=lowerCamelCase				,		focal_windows=lowerCamelCase				,		use_conv_embed=lowerCamelCase				,		idalabel=lowerCamelCase				,		labelaid=lowerCamelCase				,		use_post_layernorm=lowerCamelCase				,		use_layerscale=lowerCamelCase				,		)

	return config



def   __lowerCAmelCase					(							lowerCamelCase   :						Union[str, Any]					):






	'''simple docstring'''




	if "patch_embed.proj" in name:
		__lowerCAmelCase												=  name.replace("patch_embed.proj"				,		"embeddings.patch_embeddings.projection"					)
	if "patch_embed.norm" in name:
		__lowerCAmelCase												=  name.replace("patch_embed.norm"				,		"embeddings.norm"					)
	if "layers" in name:
		__lowerCAmelCase												=  "encoder." + name
	if "encoder.layers" in name:
		__lowerCAmelCase												=  name.replace("encoder.layers"				,		"encoder.stages"					)
	if "downsample.proj" in name:
		__lowerCAmelCase												=  name.replace("downsample.proj"				,		"downsample.projection"					)
	if "blocks" in name:
		__lowerCAmelCase												=  name.replace("blocks"				,		"layers"					)
	if "modulation.f.weight" in name or "modulation.f.bias" in name:
		__lowerCAmelCase												=  name.replace("modulation.f"				,		"modulation.projection_in"					)
	if "modulation.h.weight" in name or "modulation.h.bias" in name:
		__lowerCAmelCase												=  name.replace("modulation.h"				,		"modulation.projection_context"					)
	if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
		__lowerCAmelCase												=  name.replace("modulation.proj"				,		"modulation.projection_out"					)

	if name == "norm.weight":
		__lowerCAmelCase												=  "layernorm.weight"
	if name == "norm.bias":
		__lowerCAmelCase												=  "layernorm.bias"

	if "head" in name:
		__lowerCAmelCase												=  name.replace("head"				,		"classifier"					)
	else:
		__lowerCAmelCase												=  "focalnet." + name

	return name



def   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Any				,		lowerCamelCase   :						Union[str, Any]=False					):






	'''simple docstring'''




	__lowerCAmelCase												=  {
	    "focalnet-tiny": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth",
	    "focalnet-tiny-lrf": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth",
	    "focalnet-small": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth",
	    "focalnet-small-lrf": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth",
	    "focalnet-base": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth",
	    "focalnet-base-lrf": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth",
	    "focalnet-large-lrf-fl3": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth",
	    "focalnet-large-lrf-fl4": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth",
	    "focalnet-xlarge-lrf-fl3": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth",
	    "focalnet-xlarge-lrf-fl4": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth",
	}
	# fmt: on

	__lowerCAmelCase												=  model_name_to_url[model_name]
	print("Checkpoint URL: "				,		lowerCamelCase					)
	__lowerCAmelCase												=  torch.hub.load_state_dict_from_url(lowerCamelCase				,		map_location="cpu"					)["model"]

	# rename keys
	for key in state_dict.copy().keys():
		__lowerCAmelCase												=  state_dict.pop(lowerCamelCase					)
		__lowerCAmelCase												=  val

	__lowerCAmelCase												=  get_focalnet_config(lowerCamelCase					)
	__lowerCAmelCase												=  FocalNetForImageClassification(lowerCamelCase					)
	model.eval()

	# load state dict
	model.load_state_dict(lowerCamelCase					)

	# verify conversion
	__lowerCAmelCase												=  "http://images.cocodataset.org/val2017/000000039769.jpg"

	__lowerCAmelCase												=  BitImageProcessor(
	    do_resize=lowerCamelCase				,		size={"shortest_edge": 2_56}				,		resample=PILImageResampling.BILINEAR				,		do_center_crop=lowerCamelCase				,		crop_size=2_24				,		do_normalize=lowerCamelCase				,		image_mean=lowerCamelCase				,		image_std=lowerCamelCase				,		)
	__lowerCAmelCase												=  Image.open(requests.get(lowerCamelCase				,		stream=lowerCamelCase					).raw					)
	__lowerCAmelCase												=  processor(images=lowerCamelCase				,		return_tensors="pt"					)

	__lowerCAmelCase												=  transforms.Compose(
	    [
	        transforms.Resize(2_56					),
	        transforms.CenterCrop(2_24					),
	        transforms.ToTensor(),
	        transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]				,		std=[0.2_2_9, 0.2_2_4, 0.2_2_5]					),
	    ]					)

	__lowerCAmelCase												=  image_transforms(lowerCamelCase					).unsqueeze(0					)

	# verify pixel_values
	assert torch.allclose(inputs.pixel_values				,		lowerCamelCase				,		atol=1e-4					)

	__lowerCAmelCase												=  model(**lowerCamelCase					)

	__lowerCAmelCase												=  outputs.logits.argmax(-1					).item()
	print("Predicted class:"				,		model.config.idalabel[predicted_class_idx]					)

	print("First values of logits:"				,		outputs.logits[0, :3]					)

	if model_name == "focalnet-tiny":
		__lowerCAmelCase												=  torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]					)
	elif model_name == "focalnet-tiny-lrf":
		__lowerCAmelCase												=  torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]					)
	elif model_name == "focalnet-small":
		__lowerCAmelCase												=  torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]					)
	elif model_name == "focalnet-small-lrf":
		__lowerCAmelCase												=  torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]					)
	elif model_name == "focalnet-base":
		__lowerCAmelCase												=  torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]					)
	elif model_name == "focalnet-base-lrf":
		__lowerCAmelCase												=  torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]					)
	assert torch.allclose(outputs.logits[0, :3]				,		lowerCamelCase				,		atol=1e-4					)
	print("Looks ok!"					)

	if pytorch_dump_folder_path is not None:
		print(f'''Saving model and processor of {model_name} to {pytorch_dump_folder_path}'''					)
		model.save_pretrained(lowerCamelCase					)
		processor.save_pretrained(lowerCamelCase					)

	if push_to_hub:
		print(f'''Pushing model and processor of {model_name} to the hub...'''					)
		model.push_to_hub(f'''{model_name}'''					)
		processor.push_to_hub(f'''{model_name}'''					)


if __name__ == "__main__":
	lowerCAmelCase	:		Tuple				=   argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--model_name''',
	    default='''focalnet-tiny''',
	    type=str,
	    help='''Name of the FocalNet model you\'d like to convert.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
	)
	parser.add_argument(
	    '''--push_to_hub''',
	    action='''store_true''',
	    help='''Whether to push the model and processor to the hub.''',
	)

	lowerCAmelCase	:		Optional[int]				=   parser.parse_args()
	convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)