





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

						if "tiny" in model_name or "small" in model_name or "base" in model_name:
												UpperCamelCase				:		Tuple				   =			[3, 3, 3, 3]
												if "lrf" in model_name:
																		UpperCamelCase				:		str				   =			[3, 3, 3, 3]
												else:
																		UpperCamelCase				:		Any				   =			[2, 2, 2, 2]

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

						# set label information
						UpperCamelCase				:		Optional[int]				   =			"huggingface/label-files"
						if "large" in model_name or "huge" in model_name:
												UpperCamelCase				:		Dict				   =			"imagenet-22k-id2label.json"
						else:
												UpperCamelCase				:		List[Any]				   =			"imagenet-1k-id2label.json"

						UpperCamelCase				:		int				   =			json.load(open(hf_hub_download(lowercase__							,     lowercase__							,     repo_type="""dataset"""     )							,     """r"""     )     )
						UpperCamelCase				:		Optional[Any]				   =			{int(lowercase__     ): v for k, v in idalabel.items()}
						UpperCamelCase				:		Dict				   =			{v: k for k, v in idalabel.items()}

						UpperCamelCase				:		List[str]				   =			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      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						if "patch_embed.proj" in name:
												UpperCamelCase				:		List[Any]				   =			name.replace("""patch_embed.proj"""							,     """embeddings.patch_embeddings.projection"""     )
						if "patch_embed.norm" in name:
												UpperCamelCase				:		Union[str, Any]				   =			name.replace("""patch_embed.norm"""							,     """embeddings.norm"""     )
						if "layers" in name:
												UpperCamelCase				:		str				   =			"encoder." + name
						if "encoder.layers" in name:
												UpperCamelCase				:		Tuple				   =			name.replace("""encoder.layers"""							,     """encoder.stages"""     )
						if "downsample.proj" in name:
												UpperCamelCase				:		Optional[int]				   =			name.replace("""downsample.proj"""							,     """downsample.projection"""     )
						if "blocks" in name:
												UpperCamelCase				:		Optional[int]				   =			name.replace("""blocks"""							,     """layers"""     )
						if "modulation.f.weight" in name or "modulation.f.bias" in name:
												UpperCamelCase				:		List[Any]				   =			name.replace("""modulation.f"""							,     """modulation.projection_in"""     )
						if "modulation.h.weight" in name or "modulation.h.bias" in name:
												UpperCamelCase				:		Tuple				   =			name.replace("""modulation.h"""							,     """modulation.projection_context"""     )
						if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
												UpperCamelCase				:		Union[str, Any]				   =			name.replace("""modulation.proj"""							,     """modulation.projection_out"""     )

						if name == "norm.weight":
												UpperCamelCase				:		int				   =			"layernorm.weight"
						if name == "norm.bias":
												UpperCamelCase				:		int				   =			"layernorm.bias"

						if "head" in name:
												UpperCamelCase				:		Union[str, Any]				   =			name.replace("""head"""							,     """classifier"""     )
						else:
												UpperCamelCase				:		str				   =			"focalnet." + name

						return name

def      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE=False     ):
						# fmt: off
						UpperCamelCase				:		Optional[Any]				   =			{
						    "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

						UpperCamelCase				:		Any				   =			model_name_to_url[model_name]
						print("""Checkpoint URL: """							,     lowercase__     )
						UpperCamelCase				:		List[str]				   =			torch.hub.load_state_dict_from_url(lowercase__							,     map_location="""cpu"""     )["model"]

						# rename keys
						for key in state_dict.copy().keys():
												UpperCamelCase				:		Optional[int]				   =			state_dict.pop(lowercase__     )
												UpperCamelCase				:		Dict				   =			val

						UpperCamelCase				:		Optional[Any]				   =			get_focalnet_config(lowercase__     )
						UpperCamelCase				:		Union[str, Any]				   =			FocalNetForImageClassification(lowercase__     )
						model.eval()

						# load state dict
						model.load_state_dict(lowercase__     )

						# verify conversion
						UpperCamelCase				:		Optional[Any]				   =			"http://images.cocodataset.org/val2017/000000039769.jpg"

						UpperCamelCase				:		str				   =			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__							,     )
						UpperCamelCase				:		Dict				   =			Image.open(requests.get(lowercase__							,     stream=lowercase__     ).raw     )
						UpperCamelCase				:		str				   =			processor(images=lowercase__							,     return_tensors="""pt"""     )

						UpperCamelCase				:		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]     ),
						    ]     )

						UpperCamelCase				:		Dict				   =			image_transforms(lowercase__     ).unsqueeze(0     )

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

						UpperCamelCase				:		Dict				   =			model(**lowercase__     )

						UpperCamelCase				:		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":
												UpperCamelCase				:		Dict				   =			torch.tensor([0.21_66, -0.43_68, 0.21_91]     )
						elif model_name == "focalnet-tiny-lrf":
												UpperCamelCase				:		Tuple				   =			torch.tensor([1.16_69, 0.01_25, -0.16_95]     )
						elif model_name == "focalnet-small":
												UpperCamelCase				:		Any				   =			torch.tensor([0.49_17, -0.04_30, 0.13_41]     )
						elif model_name == "focalnet-small-lrf":
												UpperCamelCase				:		int				   =			torch.tensor([-0.25_88, -0.53_42, -0.23_31]     )
						elif model_name == "focalnet-base":
												UpperCamelCase				:		Optional[Any]				   =			torch.tensor([-0.16_55, -0.40_90, -0.17_30]     )
						elif model_name == "focalnet-base-lrf":
												UpperCamelCase				:		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__":
		__magic_name__       :					str						=			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.""",
		)

		__magic_name__       :					List[str]						=			parser.parse_args()
		convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






"""simple docstring"""

import argparse
import gc
import json
import os

import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


_UpperCamelCase		     =    1_6
_UpperCamelCase		     =    3_2



def       SCREAMING_SNAKE_CASE							( lowercase__   )			->						str:
 return int(x / 2**2_0   )






class 			__a					:





       """simple docstring"""


       def __enter__(   self   ):





        """simple docstring"""

        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
        lowerCAmelCase__   :				str 			=				torch.cuda.memory_allocated()
        return self






       def __exit__(   self			,					*snake_case   ):





        """simple docstring"""

        gc.collect()
        torch.cuda.empty_cache()
        lowerCAmelCase__   :				List[str] 			=				torch.cuda.memory_allocated()
        lowerCAmelCase__   :				Optional[int] 			=				torch.cuda.max_memory_allocated()
        lowerCAmelCase__   :				List[str] 			=				bamb(self.end - self.begin   )
        lowerCAmelCase__   :				Union[str, Any] 			=				bamb(self.peak - self.begin   )
        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")




def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__ = 1_6					,					lowercase__ = "bert-base-cased"					,					lowercase__ = 3_2_0					,					lowercase__ = 1_6_0					,					)			->						str:
 lowerCAmelCase__   :				Union[str, Any] 			=				AutoTokenizer.from_pretrained(lowercase__   )
 lowerCAmelCase__   :				Optional[Any] 			=				load_dataset(
     "glue"					,					"mrpc"					,					split={"train": F"""train[:{n_train}]""", "validation": F"""validation[:{n_val}]"""}   )

 def tokenize_function(lowercase__   ):
  # max_length=None => use the model max length (it's actually the default)
  lowerCAmelCase__   :				str 			=				tokenizer(examples["sentence1"]					,					examples["sentence2"]					,					truncation=lowercase__					,					max_length=lowercase__   )
  return outputs

 # Apply the method we just defined to all the examples in all the splits of the dataset
 lowerCAmelCase__   :				str 			=				datasets.map(
     lowercase__					,					batched=lowercase__					,					remove_columns=["idx", "sentence1", "sentence2"]					,					load_from_cache_file=lowercase__   )

 # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
 # transformers library
 lowerCAmelCase__   :				int 			=				tokenized_datasets.rename_column("label"					,					"labels"   )

 def collate_fn(lowercase__   ):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  if accelerator.distributed_type == DistributedType.TPU:
   return tokenizer.pad(lowercase__					,					padding="max_length"					,					max_length=1_2_8					,					return_tensors="pt"   )
  return tokenizer.pad(lowercase__					,					padding="longest"					,					return_tensors="pt"   )

 # Instantiate dataloaders.
 lowerCAmelCase__   :				Tuple 			=				DataLoader(
     tokenized_datasets["train"]					,					shuffle=lowercase__					,					collate_fn=lowercase__					,					batch_size=lowercase__   )
 lowerCAmelCase__   :				Optional[Any] 			=				DataLoader(
     tokenized_datasets["validation"]					,					shuffle=lowercase__					,					collate_fn=lowercase__					,					batch_size=lowercase__   )

 return train_dataloader, eval_dataloader



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__   )			->						Dict:
 # Initialize accelerator
 lowerCAmelCase__   :				Optional[Any] 			=				Accelerator()

 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 lowerCAmelCase__   :				Union[str, Any] 			=				config["lr"]
 lowerCAmelCase__   :				int 			=				int(config["num_epochs"]   )
 lowerCAmelCase__   :				Tuple 			=				int(config["seed"]   )
 lowerCAmelCase__   :				str 			=				int(config["batch_size"]   )
 lowerCAmelCase__   :				Any 			=				args.model_name_or_path

 set_seed(lowercase__   )
 lowerCAmelCase__			,						lowerCAmelCase__   :				Optional[Any] 			=				get_dataloaders(lowercase__					,					lowercase__					,					lowercase__					,					args.n_train					,					args.n_val   )

 # Instantiate the model (we build the model here so that the seed also control new weights initialization)
 lowerCAmelCase__   :				Any 			=				AutoModelForSequenceClassification.from_pretrained(lowercase__					,					return_dict=lowercase__   )

 # Instantiate optimizer
 lowerCAmelCase__   :				List[str] 			=				(
     AdamW
     if accelerator.state.deepspeed_plugin is None
     or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
     else DummyOptim
 )
 lowerCAmelCase__   :				Any 			=				optimizer_cls(params=model.parameters()					,					lr=lowercase__   )

 if accelerator.state.deepspeed_plugin is not None:
  lowerCAmelCase__   :				List[str] 			=				accelerator.state.deepspeed_plugin.deepspeed_config[
      "gradient_accumulation_steps"
  ]
 else:
  lowerCAmelCase__   :				Dict 			=				1
 lowerCAmelCase__   :				List[str] 			=				(len(lowercase__   ) * num_epochs) // gradient_accumulation_steps

 # Instantiate scheduler
 if (
     accelerator.state.deepspeed_plugin is None
     or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
 ):
  lowerCAmelCase__   :				Any 			=				get_linear_schedule_with_warmup(
      optimizer=lowercase__					,					num_warmup_steps=0					,					num_training_steps=lowercase__					,					)
 else:
  lowerCAmelCase__   :				Optional[int] 			=				DummyScheduler(lowercase__					,					total_num_steps=lowercase__					,					warmup_num_steps=0   )

 # Prepare everything
 # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
 # prepare method.
 lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Dict 			=				accelerator.prepare(
     lowercase__					,					lowercase__					,					lowercase__					,					lowercase__					,					lowercase__   )

 # We need to keep track of how many total steps we have iterated over
 lowerCAmelCase__   :				str 			=				0
 # We also need to keep track of the stating epoch so files are named properly
 lowerCAmelCase__   :				Tuple 			=				0

 # Now we train the model
 lowerCAmelCase__   :				List[Any] 			=				{}
 for epoch in range(lowercase__					,					lowercase__   ):
  with TorchTracemalloc() as tracemalloc:
   model.train()
   for step, batch in enumerate(lowercase__   ):
    lowerCAmelCase__   :				Optional[Any] 			=				model(**lowercase__   )
    lowerCAmelCase__   :				Optional[int] 			=				outputs.loss
    lowerCAmelCase__   :				Any 			=				loss / gradient_accumulation_steps
    accelerator.backward(lowercase__   )
    if step % gradient_accumulation_steps == 0:
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()

    overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
  accelerator.print("Memory before entering the train : {}".format(bamb(tracemalloc.begin   )   )   )
  accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used   )   )
  accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked   )   )
  accelerator.print(
      "Total Peak Memory consumed during the train (max): {}".format(
          tracemalloc.peaked + bamb(tracemalloc.begin   )   )   )
  lowerCAmelCase__   :				Optional[int] 			=				tracemalloc.peaked + bamb(tracemalloc.begin   )
  if args.peak_memory_upper_bound is not None:
   assert (
       train_total_peak_memory[F"""epoch-{epoch}"""] <= args.peak_memory_upper_bound
   ), "Peak memory usage exceeded the upper bound"

 accelerator.wait_for_everyone()
 if accelerator.is_main_process:
  with open(os.path.join(args.output_dir					,					"peak_memory_utilization.json"   )					,					"w"   ) as f:
   json.dump(lowercase__					,					lowercase__   )



def       SCREAMING_SNAKE_CASE							( )			->						Optional[Any]:
 lowerCAmelCase__   :				int 			=				argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage."   )
 parser.add_argument(
     "--model_name_or_path"					,					type=lowercase__					,					default="bert-base-cased"					,					help="Path to pretrained model or model identifier from huggingface.co/models."					,					required=lowercase__					,					)
 parser.add_argument(
     "--output_dir"					,					type=lowercase__					,					default="."					,					help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory."					,					)
 parser.add_argument(
     "--peak_memory_upper_bound"					,					type=lowercase__					,					default=lowercase__					,					help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value."					,					)
 parser.add_argument(
     "--n_train"					,					type=lowercase__					,					default=3_2_0					,					help="Number of training examples to use."					,					)
 parser.add_argument(
     "--n_val"					,					type=lowercase__					,					default=1_6_0					,					help="Number of validation examples to use."					,					)
 parser.add_argument(
     "--num_epochs"					,					type=lowercase__					,					default=1					,					help="Number of train epochs."					,					)
 lowerCAmelCase__   :				Optional[int] 			=				parser.parse_args()
 lowerCAmelCase__   :				Dict 			=				{"lr": 2E-5, "num_epochs": args.num_epochs, "seed": 4_2, "batch_size": 1_6}
 training_function(lowercase__					,					lowercase__   )


if __name__ == "__main__":
    main()






