






from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
	import tensorflow as tf

	from transformers.models.xglm.modeling_tf_xglm import (
	    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFXGLMForCausalLM,
	    TFXGLMModel,
	)

@require_tf
class     _SCREAMING_SNAKE_CASE       :
	snake_case__							:     List[str]     	=  XGLMConfig
	snake_case__							:     Optional[Any]     	=  {}
	snake_case__							:     Optional[int]     	=  """gelu"""





	def __init__(		self				:					Union[str, Any] ,  __lowerCamelCase				:					str ,  __lowerCamelCase				:					Optional[Any]=14 ,  __lowerCamelCase				:					Union[str, Any]=7 ,  __lowerCamelCase				:					List[Any]=True ,  __lowerCamelCase				:					Any=True ,  __lowerCamelCase				:					int=True ,  __lowerCamelCase				:					Any=99 ,  __lowerCamelCase				:					Union[str, Any]=32 ,  __lowerCamelCase				:					List[Any]=2 ,  __lowerCamelCase				:					Optional[Any]=4 ,  __lowerCamelCase				:					Optional[Any]=37 ,  __lowerCamelCase				:					int="gelu" ,  __lowerCamelCase				:					int=0.1 ,  __lowerCamelCase				:					Any=0.1 ,  __lowerCamelCase				:					List[str]=512 ,  __lowerCamelCase				:					List[Any]=0.02 ,  ):
		UpperCamelCase     :Tuple          =	parent
		UpperCamelCase     :Optional[Any]          =	batch_size
		UpperCamelCase     :List[Any]          =	seq_length
		UpperCamelCase     :Union[str, Any]          =	is_training
		UpperCamelCase     :Tuple          =	use_input_mask
		UpperCamelCase     :List[str]          =	use_labels
		UpperCamelCase     :Union[str, Any]          =	vocab_size
		UpperCamelCase     :Tuple          =	d_model
		UpperCamelCase     :Any          =	num_hidden_layers
		UpperCamelCase     :Union[str, Any]          =	num_attention_heads
		UpperCamelCase     :Optional[Any]          =	ffn_dim
		UpperCamelCase     :List[Any]          =	activation_function
		UpperCamelCase     :str          =	activation_dropout
		UpperCamelCase     :List[str]          =	attention_dropout
		UpperCamelCase     :List[str]          =	max_position_embeddings
		UpperCamelCase     :Dict          =	initializer_range
		UpperCamelCase     :List[Any]          =	None
		UpperCamelCase     :str          =	0
		UpperCamelCase     :Any          =	2
		UpperCamelCase     :Any          =	1





	def 						_A (		self				:					int	):
		return XGLMConfig.from_pretrained("""facebook/xglm-564M"""	)





	def 						_A (		self				:					List[str]	):
		UpperCamelCase     :Union[str, Any]          =	tf.clip_by_value(
		    ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size	) ,  clip_value_min=0 ,  clip_value_max=3	)

		UpperCamelCase     :List[str]          =	None
		if self.use_input_mask:
			UpperCamelCase     :List[Any]          =	random_attention_mask([self.batch_size, self.seq_length]	)

		UpperCamelCase     :str          =	self.get_config()

		UpperCamelCase     :List[Any]          =	floats_tensor([self.num_hidden_layers, self.num_attention_heads] ,  2	)

		return (
		    config,
		    input_ids,
		    input_mask,
		    head_mask,
		)





	def 						_A (		self				:					str	):
		return XGLMConfig(
		    vocab_size=self.vocab_size ,  d_model=self.hidden_size ,  num_layers=self.num_hidden_layers ,  attention_heads=self.num_attention_heads ,  ffn_dim=self.ffn_dim ,  activation_function=self.activation_function ,  activation_dropout=self.activation_dropout ,  attention_dropout=self.attention_dropout ,  max_position_embeddings=self.max_position_embeddings ,  initializer_range=self.initializer_range ,  use_cache=__lowerCamelCase ,  bos_token_id=self.bos_token_id ,  eos_token_id=self.eos_token_id ,  pad_token_id=self.pad_token_id ,  return_dict=__lowerCamelCase ,  )





	def 						_A (		self				:					Union[str, Any]	):
		UpperCamelCase     :Optional[Any]          =	self.prepare_config_and_inputs()

		(
		    (
		    UpperCamelCase
		)     ,      (
		    UpperCamelCase
		)     ,      (
		    UpperCamelCase
		)     ,      (
		    UpperCamelCase
		)     ,      
		)     :Optional[Any]          =	config_and_inputs

		UpperCamelCase     :str          =	{
		    """input_ids""": input_ids,
		    """head_mask""": head_mask,
		}

		return config, inputs_dict

@require_tf
class     _SCREAMING_SNAKE_CASE       (			_a		,   _a		,   unittest.TestCase      ):
	snake_case__							:     int     	=  (TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
	snake_case__							:     Dict     	=  (TFXGLMForCausalLM,) if is_tf_available() else ()
	snake_case__							:     Optional[Any]     	=  (
	    {"""feature-extraction""": TFXGLMModel, """text-generation""": TFXGLMForCausalLM} if is_tf_available() else {}
	)
	snake_case__							:     List[str]     	=  False
	snake_case__							:     Optional[int]     	=  False
	snake_case__							:     List[str]     	=  False





	def 						_A (		self				:					List[Any]	):
		UpperCamelCase     :Optional[Any]          =	TFXGLMModelTester(self	)
		UpperCamelCase     :List[Any]          =	ConfigTester(self ,  config_class=__lowerCamelCase ,  n_embd=37	)





	def 						_A (		self				:					int	):
		self.config_tester.run_common_tests()





	@slow
	def 						_A (		self				:					int	):
		for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
			UpperCamelCase     :str          =	TFXGLMModel.from_pretrained(__lowerCamelCase	)
			self.assertIsNotNone(__lowerCamelCase	)





	@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""	)
	def 						_A (		self				:					Dict	):
		super().test_resize_token_embeddings()

@require_tf
class     _SCREAMING_SNAKE_CASE       (			unittest.TestCase      ):





	@slow
	def 						_A (		self				:					Optional[Any] ,  __lowerCamelCase				:					Optional[int]=True	):
		UpperCamelCase     :Dict          =	TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""	)
		UpperCamelCase     :str          =	tf.convert_to_tensor([[2, 268, 9_865]] ,  dtype=tf.intaa	)  # The dog
		# </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
		# fmt: off
		UpperCamelCase     :int          =	[2, 268, 9_865, 67, 11, 1_988, 57_252, 9_865, 5, 984, 67, 1_988, 213_838, 1_658, 53, 70_446, 33, 6_657, 278, 1_581]
		# fmt: on
		UpperCamelCase     :str          =	model.generate(__lowerCamelCase ,  do_sample=__lowerCamelCase ,  num_beams=1	)
		if verify_outputs:
			self.assertListEqual(output_ids[0].numpy().tolist() ,  __lowerCamelCase	)





	@slow
	def 						_A (		self				:					Tuple	):
		UpperCamelCase     :Tuple          =	XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""	)
		UpperCamelCase     :Tuple          =	TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""	)

		tf.random.set_seed(0	)
		UpperCamelCase     :Optional[int]          =	tokenizer("""Today is a nice day and""" ,  return_tensors="""tf"""	)
		UpperCamelCase     :str          =	tokenized.input_ids
		# forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
		with tf.device(""":/CPU:0"""	):
			UpperCamelCase     :Any          =	model.generate(__lowerCamelCase ,  do_sample=__lowerCamelCase ,  seed=[7, 0]	)
		UpperCamelCase     :str          =	tokenizer.decode(output_ids[0] ,  skip_special_tokens=__lowerCamelCase	)

		UpperCamelCase     :Union[str, Any]          =	(
		    """Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due"""
		)
		self.assertEqual(__lowerCamelCase ,  __lowerCamelCase	)





	@slow
	def 						_A (		self				:					int	):
		UpperCamelCase     :List[str]          =	TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""	)
		UpperCamelCase     :Tuple          =	XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""	)

		UpperCamelCase     :str          =	"""left"""

		# use different length sentences to test batching
		UpperCamelCase     :Union[str, Any]          =	[
		    """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
		    """left-padding, such as in batched generation. The output for the sequence below should be the same """
		    """regardless of whether left padding is applied or not. When""",
		    """Hello, my dog is a little""",
		]

		UpperCamelCase     :str          =	tokenizer(__lowerCamelCase ,  return_tensors="""tf""" ,  padding=__lowerCamelCase	)
		UpperCamelCase     :Optional[int]          =	inputs["""input_ids"""]

		UpperCamelCase     :Dict          =	model.generate(input_ids=__lowerCamelCase ,  attention_mask=inputs["""attention_mask"""] ,  max_new_tokens=12	)

		UpperCamelCase     :List[str]          =	tokenizer(sentences[0] ,  return_tensors="""tf"""	).input_ids
		UpperCamelCase     :List[Any]          =	model.generate(input_ids=__lowerCamelCase ,  max_new_tokens=12	)

		UpperCamelCase     :str          =	tokenizer(sentences[1] ,  return_tensors="""tf"""	).input_ids
		UpperCamelCase     :str          =	model.generate(input_ids=__lowerCamelCase ,  max_new_tokens=12	)

		UpperCamelCase     :Any          =	tokenizer.batch_decode(__lowerCamelCase ,  skip_special_tokens=__lowerCamelCase	)
		UpperCamelCase     :Union[str, Any]          =	tokenizer.decode(output_non_padded[0] ,  skip_special_tokens=__lowerCamelCase	)
		UpperCamelCase     :List[Any]          =	tokenizer.decode(output_padded[0] ,  skip_special_tokens=__lowerCamelCase	)

		UpperCamelCase     :List[Any]          =	[
		    """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
		    """left-padding, such as in batched generation. The output for the sequence below should be the same """
		    """regardless of whether left padding is applied or not. When left padding is applied, the sequence will be """
		    """a single""",
		    """Hello, my dog is a little bit of a shy one, but he is very friendly""",
		]
		self.assertListEqual(__lowerCamelCase ,  __lowerCamelCase	)
		self.assertListEqual(__lowerCamelCase ,  [non_padded_sentence, padded_sentence]	)






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 __UpperCamelCase			(       _A  ):
						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_      			=  128
						elif "large" in model_name:
												lowerCAmelCase_      			=  192
						elif "xlarge" in model_name:
												lowerCAmelCase_      			=  256
						elif "huge" in model_name:
												lowerCAmelCase_      			=  352

						# 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(_A ,  _A ,  repo_type='''dataset'''  ) ,  '''r'''  )  )
						lowerCAmelCase_      			=  {int(_A  ): v for k, v in idalabel.items()}
						lowerCAmelCase_      			=  {v: k for k, v in idalabel.items()}

						lowerCAmelCase_      			=  FocalNetConfig(
						    embed_dim=_A ,  depths=_A ,  focal_levels=_A ,  focal_windows=_A ,  use_conv_embed=_A ,  idalabel=_A ,  labelaid=_A ,  use_post_layernorm=_A ,  use_layerscale=_A ,  )

						return config





def __UpperCamelCase			(       _A  ):
						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 __UpperCamelCase			(       _A ,  _A ,  _A=False  ):
						# fmt: off
						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: ''' ,  _A  )
						lowerCAmelCase_      			=  torch.hub.load_state_dict_from_url(_A ,  map_location='''cpu'''  )['''model''']

						# rename keys
						for key in state_dict.copy().keys():
												lowerCAmelCase_      			=  state_dict.pop(_A  )
												lowerCAmelCase_      			=  val

						lowerCAmelCase_      			=  get_focalnet_config(_A  )
						lowerCAmelCase_      			=  FocalNetForImageClassification(_A  )
						model.eval()

						# load state dict
						model.load_state_dict(_A  )

						# verify conversion
						lowerCAmelCase_      			=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''

						lowerCAmelCase_      			=  BitImageProcessor(
						    do_resize=_A ,  size={'''shortest_edge''': 256} ,  resample=PILImageResampling.BILINEAR ,  do_center_crop=_A ,  crop_size=224 ,  do_normalize=_A ,  image_mean=_A ,  image_std=_A ,  )
						lowerCAmelCase_      			=  Image.open(requests.get(_A ,  stream=_A  ).raw  )
						lowerCAmelCase_      			=  processor(images=_A ,  return_tensors='''pt'''  )

						lowerCAmelCase_      			=  transforms.Compose(
						    [
						        transforms.Resize(256  ),
						        transforms.CenterCrop(224  ),
						        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(_A  ).unsqueeze(0  )

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

						lowerCAmelCase_      			=  model(**_A  )

						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] ,  _A ,  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(_A  )
												processor.save_pretrained(_A  )

						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__":
						_A     =			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.''',
						)

						_A     =			parser.parse_args()
						convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




