
'''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 lowercase_	(   _lowercase				)     ->							Dict:
					'''simple docstring'''

					lowerCamelCase_					:      Optional[int]     =		[2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
					lowerCamelCase_					:      List[str]     =		True if '''large''' in model_name or '''huge''' in model_name else False
					lowerCamelCase_					:      Optional[int]     =		True if '''large''' in model_name or '''huge''' in model_name else False
					lowerCamelCase_					:      Any     =		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_					:      int     =		[3, 3, 3, 3]
															lowerCamelCase_					:      List[Any]     =		[5, 5, 5, 5]
										elif "fl4" in model_name:
															lowerCamelCase_					:      str     =		[4, 4, 4, 4]
															lowerCamelCase_					:      int     =		[3, 3, 3, 3]

					if "tiny" in model_name or "small" in model_name or "base" in model_name:
										lowerCamelCase_					:      Optional[int]     =		[3, 3, 3, 3]
										if "lrf" in model_name:
															lowerCamelCase_					:      str     =		[3, 3, 3, 3]
										else:
															lowerCamelCase_					:      Union[str, Any]     =		[2, 2, 2, 2]

					if "tiny" in model_name:
										lowerCamelCase_					:      Optional[int]     =		96
					elif "small" in model_name:
										lowerCamelCase_					:      Dict     =		96
					elif "base" in model_name:
										lowerCamelCase_					:      Tuple     =		128
					elif "large" in model_name:
										lowerCamelCase_					:      List[str]     =		192
					elif "xlarge" in model_name:
										lowerCamelCase_					:      Any     =		256
					elif "huge" in model_name:
										lowerCamelCase_					:      Union[str, Any]     =		352

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

					lowerCamelCase_					:      str     =		json.load(open(hf_hub_download(_lowercase			,       _lowercase			,       repo_type='''dataset'''				)			,       '''r'''				)				)
					lowerCamelCase_					:      int     =		{int(_lowercase				): v for k, v in idalabel.items()}
					lowerCamelCase_					:      int     =		{v: k for k, v in idalabel.items()}

					lowerCamelCase_					:      int     =		FocalNetConfig(
					    embed_dim=_lowercase			,       depths=_lowercase			,       focal_levels=_lowercase			,       focal_windows=_lowercase			,       use_conv_embed=_lowercase			,       idalabel=_lowercase			,       labelaid=_lowercase			,       use_post_layernorm=_lowercase			,       use_layerscale=_lowercase			,       )

					return config




def lowercase_	(   _lowercase				)     ->							Union[str, Any]:
					'''simple docstring'''

					if "patch_embed.proj" in name:
										lowerCamelCase_					:      Union[str, Any]     =		name.replace('''patch_embed.proj'''			,       '''embeddings.patch_embeddings.projection'''				)
					if "patch_embed.norm" in name:
										lowerCamelCase_					:      Any     =		name.replace('''patch_embed.norm'''			,       '''embeddings.norm'''				)
					if "layers" in name:
										lowerCamelCase_					:      Tuple     =		'''encoder.''' + name
					if "encoder.layers" in name:
										lowerCamelCase_					:      Optional[Any]     =		name.replace('''encoder.layers'''			,       '''encoder.stages'''				)
					if "downsample.proj" in name:
										lowerCamelCase_					:      Any     =		name.replace('''downsample.proj'''			,       '''downsample.projection'''				)
					if "blocks" in name:
										lowerCamelCase_					:      List[Any]     =		name.replace('''blocks'''			,       '''layers'''				)
					if "modulation.f.weight" in name or "modulation.f.bias" in name:
										lowerCamelCase_					:      Union[str, Any]     =		name.replace('''modulation.f'''			,       '''modulation.projection_in'''				)
					if "modulation.h.weight" in name or "modulation.h.bias" in name:
										lowerCamelCase_					:      str     =		name.replace('''modulation.h'''			,       '''modulation.projection_context'''				)
					if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
										lowerCamelCase_					:      Any     =		name.replace('''modulation.proj'''			,       '''modulation.projection_out'''				)

					if name == "norm.weight":
										lowerCamelCase_					:      Optional[Any]     =		'''layernorm.weight'''
					if name == "norm.bias":
										lowerCamelCase_					:      Optional[Any]     =		'''layernorm.bias'''

					if "head" in name:
										lowerCamelCase_					:      Optional[Any]     =		name.replace('''head'''			,       '''classifier'''				)
					else:
										lowerCamelCase_					:      List[Any]     =		'''focalnet.''' + name

					return name




def lowercase_	(   _lowercase			,       _lowercase			,       _lowercase=False				)     ->							int:
					'''simple docstring'''

					lowerCamelCase_					:      int     =		{
					    '''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_					:      Tuple     =		model_name_to_url[model_name]
					print('''Checkpoint URL: '''			,       _lowercase				)
					lowerCamelCase_					:      Dict     =		torch.hub.load_state_dict_from_url(_lowercase			,       map_location='''cpu'''				)['''model''']

					# rename keys
					for key in state_dict.copy().keys():
										lowerCamelCase_					:      Any     =		state_dict.pop(_lowercase				)
										lowerCamelCase_					:      Union[str, Any]     =		val

					lowerCamelCase_					:      Any     =		get_focalnet_config(_lowercase				)
					lowerCamelCase_					:      Tuple     =		FocalNetForImageClassification(_lowercase				)
					model.eval()

					# load state dict
					model.load_state_dict(_lowercase				)

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

					lowerCamelCase_					:      Optional[int]     =		BitImageProcessor(
					    do_resize=_lowercase			,       size={'''shortest_edge''': 256}			,       resample=PILImageResampling.BILINEAR			,       do_center_crop=_lowercase			,       crop_size=224			,       do_normalize=_lowercase			,       image_mean=_lowercase			,       image_std=_lowercase			,       )
					lowerCamelCase_					:      Optional[Any]     =		Image.open(requests.get(_lowercase			,       stream=_lowercase				).raw				)
					lowerCamelCase_					:      Optional[Any]     =		processor(images=_lowercase			,       return_tensors='''pt'''				)

					lowerCamelCase_					:      Tuple     =		transforms.Compose(
					    [
					        transforms.Resize(256				),
					        transforms.CenterCrop(224				),
					        transforms.ToTensor(),
					        transforms.Normalize(mean=[0.4_85, 0.4_56, 0.4_06]			,       std=[0.2_29, 0.2_24, 0.2_25]				),
					    ]				)

					lowerCamelCase_					:      Dict     =		image_transforms(_lowercase				).unsqueeze(0				)

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

					lowerCamelCase_					:      Dict     =		model(**_lowercase				)

					lowerCamelCase_					:      Dict     =		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_					:      List[Any]     =		torch.tensor([0.21_66, -0.43_68, 0.21_91]				)
					elif model_name == "focalnet-tiny-lrf":
										lowerCamelCase_					:      str     =		torch.tensor([1.16_69, 0.01_25, -0.16_95]				)
					elif model_name == "focalnet-small":
										lowerCamelCase_					:      Dict     =		torch.tensor([0.49_17, -0.04_30, 0.13_41]				)
					elif model_name == "focalnet-small-lrf":
										lowerCamelCase_					:      List[Any]     =		torch.tensor([-0.25_88, -0.53_42, -0.23_31]				)
					elif model_name == "focalnet-base":
										lowerCamelCase_					:      Union[str, Any]     =		torch.tensor([-0.16_55, -0.40_90, -0.17_30]				)
					elif model_name == "focalnet-base-lrf":
										lowerCamelCase_					:      Tuple     =		torch.tensor([0.53_06, -0.04_83, -0.39_28]				)
					assert torch.allclose(outputs.logits[0, :3]			,       _lowercase			,       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(_lowercase				)
										processor.save_pretrained(_lowercase				)

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

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




'''simple docstring'''


import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer

class 	__lowercase				(  _lowercase    ):

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Any     =		tempfile.mkdtemp()
								lowerCamelCase_					:      Optional[int]     =		5

								# Realm tok
								lowerCamelCase_					:      Any     =		[
								    '''[UNK]''',
								    '''[CLS]''',
								    '''[SEP]''',
								    '''[PAD]''',
								    '''[MASK]''',
								    '''test''',
								    '''question''',
								    '''this''',
								    '''is''',
								    '''the''',
								    '''first''',
								    '''second''',
								    '''third''',
								    '''fourth''',
								    '''fifth''',
								    '''record''',
								    '''want''',
								    '''##want''',
								    '''##ed''',
								    '''wa''',
								    '''un''',
								    '''runn''',
								    '''##ing''',
								    ''',''',
								    '''low''',
								    '''lowest''',
								]
								lowerCamelCase_					:      Optional[int]     =		os.path.join(self.tmpdirname			,   '''realm_tokenizer'''      )
								os.makedirs(A			,   exist_ok=A      )
								lowerCamelCase_					:      str     =		os.path.join(A			,   VOCAB_FILES_NAMES['''vocab_file''']      )
								with open(self.vocab_file			,   '''w'''			,   encoding='''utf-8'''      ) as vocab_writer:
													vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]      )      )

								lowerCamelCase_					:      List[Any]     =		os.path.join(self.tmpdirname			,   '''realm_block_records'''      )
								os.makedirs(A			,   exist_ok=A      )

			def        UpperCAmelCase__ (self      ):
								return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,   '''realm_tokenizer'''      )      )

			def        UpperCAmelCase__ (self      ):
								shutil.rmtree(self.tmpdirname      )

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      int     =		RealmConfig(num_block_records=self.num_block_records      )
								return config

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[int]     =		Dataset.from_dict(
								    {
								        '''id''': ['''0''', '''1'''],
								        '''question''': ['''foo''', '''bar'''],
								        '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
								    }      )
								return dataset

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Any     =		np.array(
								    [
								        B'''This is the first record''',
								        B'''This is the second record''',
								        B'''This is the third record''',
								        B'''This is the fourth record''',
								        B'''This is the fifth record''',
								        B'''This is a longer longer longer record''',
								    ]			,   dtype=A			,   )
								return block_records

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Dict     =		RealmRetriever(
								    block_records=self.get_dummy_block_records()			,   tokenizer=self.get_tokenizer()			,   )
								return retriever

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[Any]     =		self.get_config()
								lowerCamelCase_					:      List[Any]     =		self.get_dummy_retriever()
								lowerCamelCase_					:      Dict     =		retriever.tokenizer

								lowerCamelCase_					:      List[Any]     =		np.array([0, 3]			,   dtype='''long'''      )
								lowerCamelCase_					:      int     =		tokenizer(['''Test question''']      ).input_ids
								lowerCamelCase_					:      Any     =		tokenizer(
								    ['''the fourth''']			,   add_special_tokens=A			,   return_token_type_ids=A			,   return_attention_mask=A			,   ).input_ids
								lowerCamelCase_					:      Any     =		config.reader_seq_len

								lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_					:      Tuple     =		retriever(
								    A			,   A			,   answer_ids=A			,   max_length=A			,   return_tensors='''np'''      )

								self.assertEqual(len(A      )			,   2      )
								self.assertEqual(len(A      )			,   2      )
								self.assertEqual(len(A      )			,   2      )
								self.assertEqual(concat_inputs.input_ids.shape			,   (2, 1_0)      )
								self.assertEqual(concat_inputs.attention_mask.shape			,   (2, 1_0)      )
								self.assertEqual(concat_inputs.token_type_ids.shape			,   (2, 1_0)      )
								self.assertEqual(concat_inputs.special_tokens_mask.shape			,   (2, 1_0)      )
								self.assertEqual(
								    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]      )			,   ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''first''', '''record''', '''[SEP]''']			,   )
								self.assertEqual(
								    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]      )			,   ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''fourth''', '''record''', '''[SEP]''']			,   )

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      str     =		self.get_config()
								lowerCamelCase_					:      Optional[Any]     =		self.get_dummy_retriever()
								lowerCamelCase_					:      Optional[int]     =		retriever.tokenizer

								lowerCamelCase_					:      int     =		np.array([0, 3, 5]			,   dtype='''long'''      )
								lowerCamelCase_					:      List[Any]     =		tokenizer(['''Test question''']      ).input_ids
								lowerCamelCase_					:      int     =		tokenizer(
								    ['''the fourth''', '''longer longer''']			,   add_special_tokens=A			,   return_token_type_ids=A			,   return_attention_mask=A			,   ).input_ids
								lowerCamelCase_					:      Tuple     =		config.reader_seq_len

								lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_					:      int     =		retriever(
								    A			,   A			,   answer_ids=A			,   max_length=A			,   return_tensors='''np'''      )

								self.assertEqual([False, True, True]			,   A      )
								self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]			,   A      )
								self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]			,   A      )






			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Any     =		self.get_dummy_retriever()
								retriever.save_pretrained(os.path.join(self.tmpdirname			,   '''realm_block_records'''      )      )

								# Test local path
								lowerCamelCase_					:      str     =		retriever.from_pretrained(os.path.join(self.tmpdirname			,   '''realm_block_records'''      )      )
								self.assertEqual(retriever.block_records[0]			,   B'''This is the first record'''      )

								# Test mocked remote path
								with patch('''transformers.models.realm.retrieval_realm.hf_hub_download'''      ) as mock_hf_hub_download:
													lowerCamelCase_					:      Dict     =		os.path.join(
													    os.path.join(self.tmpdirname			,   '''realm_block_records'''      )			,   _REALM_BLOCK_RECORDS_FILENAME      )
													lowerCamelCase_					:      str     =		RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''      )

								self.assertEqual(retriever.block_records[0]			,   B'''This is the first record'''      )


