





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 							a_	(    __magic_name__   )				->			List[str]:




				"""simple docstring"""






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

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

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

				# set label information
				snake_case   :					int    						=				'''huggingface/label-files'''
				if "large" in model_name or "huge" in model_name:
								snake_case   :					Tuple    						=				'''imagenet-22k-id2label.json'''
				else:
								snake_case   :					List[str]    						=				'''imagenet-1k-id2label.json'''

				snake_case   :					Optional[int]    						=				json.load(open(hf_hub_download(__magic_name__		,     __magic_name__		,     repo_type='''dataset'''   )		,     '''r'''   )   )
				snake_case   :					List[Any]    						=				{int(__magic_name__   ): v for k, v in idalabel.items()}
				snake_case   :					Union[str, Any]    						=				{v: k for k, v in idalabel.items()}

				snake_case   :					Optional[Any]    						=				FocalNetConfig(
				    embed_dim=__magic_name__		,     depths=__magic_name__		,     focal_levels=__magic_name__		,     focal_windows=__magic_name__		,     use_conv_embed=__magic_name__		,     idalabel=__magic_name__		,     labelaid=__magic_name__		,     use_post_layernorm=__magic_name__		,     use_layerscale=__magic_name__		,     )

				return config
def 							a_	(    __magic_name__   )				->			Dict:




				"""simple docstring"""






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

				if name == "norm.weight":
								snake_case   :					Optional[int]    						=				'''layernorm.weight'''
				if name == "norm.bias":
								snake_case   :					Optional[int]    						=				'''layernorm.bias'''

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

				return name
def 							a_	(    __magic_name__		,     __magic_name__		,     __magic_name__=False   )				->			Union[str, Any]:




				"""simple docstring"""






				snake_case   :					str    						=				{
				    '''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

				snake_case   :					List[Any]    						=				model_name_to_url[model_name]
				print('''Checkpoint URL: '''		,     __magic_name__   )
				snake_case   :					List[str]    						=				torch.hub.load_state_dict_from_url(__magic_name__		,     map_location='''cpu'''   )['''model''']

				# rename keys
				for key in state_dict.copy().keys():
								snake_case   :					List[Any]    						=				state_dict.pop(__magic_name__   )
								snake_case   :					Any    						=				val

				snake_case   :					Union[str, Any]    						=				get_focalnet_config(__magic_name__   )
				snake_case   :					Dict    						=				FocalNetForImageClassification(__magic_name__   )
				model.eval()

				# load state dict
				model.load_state_dict(__magic_name__   )

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

				snake_case   :					Union[str, Any]    						=				BitImageProcessor(
				    do_resize=__magic_name__		,     size={'''shortest_edge''': 256}		,     resample=PILImageResampling.BILINEAR		,     do_center_crop=__magic_name__		,     crop_size=224		,     do_normalize=__magic_name__		,     image_mean=__magic_name__		,     image_std=__magic_name__		,     )
				snake_case   :					List[str]    						=				Image.open(requests.get(__magic_name__		,     stream=__magic_name__   ).raw   )
				snake_case   :					Dict    						=				processor(images=__magic_name__		,     return_tensors='''pt'''   )

				snake_case   :					Optional[Any]    						=				transforms.Compose(
				    [
				        transforms.Resize(256   ),
				        transforms.CenterCrop(224   ),
				        transforms.ToTensor(),
				        transforms.Normalize(mean=[0.485, 0.456, 0.406]		,     std=[0.229, 0.224, 0.225]   ),
				    ]   )

				snake_case   :					int    						=				image_transforms(__magic_name__   ).unsqueeze(0   )

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

				snake_case   :					int    						=				model(**__magic_name__   )

				snake_case   :					Any    						=				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":
								snake_case   :					str    						=				torch.tensor([0.2166, -0.4368, 0.2191]   )
				elif model_name == "focalnet-tiny-lrf":
								snake_case   :					Optional[int]    						=				torch.tensor([1.1669, 0.0125, -0.1695]   )
				elif model_name == "focalnet-small":
								snake_case   :					Optional[int]    						=				torch.tensor([0.4917, -0.0430, 0.1341]   )
				elif model_name == "focalnet-small-lrf":
								snake_case   :					str    						=				torch.tensor([-0.2588, -0.5342, -0.2331]   )
				elif model_name == "focalnet-base":
								snake_case   :					List[Any]    						=				torch.tensor([-0.1655, -0.4090, -0.1730]   )
				elif model_name == "focalnet-base-lrf":
								snake_case   :					int    						=				torch.tensor([0.5306, -0.0483, -0.3928]   )
				assert torch.allclose(outputs.logits[0, :3]		,     __magic_name__		,     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(__magic_name__   )
								processor.save_pretrained(__magic_name__   )

				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     :    List[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.',
			)

			_a     :    Any         =						parser.parse_args()
			convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






'''simple docstring'''







import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


lowerCAmelCase_     						=				1_6
lowerCAmelCase_     						=				3_2







def 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase = 16						):






   '''simple docstring'''







   A__       =     AutoTokenizer.from_pretrained('bert-base-cased'						)
   A__       =     DatasetDict(
       {
           'train': dataset['train'].select(UpperCAmelCase						),
           'validation': dataset['train'].select(UpperCAmelCase						),
           'test': dataset['validation'],
       }						)

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

   # Apply the method we just defined to all the examples in all the splits of the dataset
   # starting with the main process first:
   with accelerator.main_process_first():
      A__       =     datasets.map(
          UpperCAmelCase      ,batched=UpperCAmelCase      ,remove_columns=['idx', 'sentence1', 'sentence2']      ,)

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

   def collate_fn(UpperCAmelCase						):
      # On TPU it's best to pad everything to the same length or training will be very slow.
      A__       =     128 if accelerator.distributed_type == DistributedType.TPU else None
      # When using mixed precision we want round multiples of 8/16
      if accelerator.mixed_precision == "fp8":
         A__       =     16
      elif accelerator.mixed_precision != "no":
         A__       =     8
      else:
         A__       =     None

      return tokenizer.pad(
          UpperCAmelCase      ,padding='longest'      ,max_length=UpperCAmelCase      ,pad_to_multiple_of=UpperCAmelCase      ,return_tensors='pt'      ,)

   # Instantiate dataloaders.
   A__       =     DataLoader(
       tokenized_datasets['train']      ,shuffle=UpperCAmelCase      ,collate_fn=UpperCAmelCase      ,batch_size=UpperCAmelCase						)
   A__       =     DataLoader(
       tokenized_datasets['validation']      ,shuffle=UpperCAmelCase      ,collate_fn=UpperCAmelCase      ,batch_size=UpperCAmelCase						)

   A__       =     DataLoader(
       tokenized_datasets['test']      ,shuffle=UpperCAmelCase      ,collate_fn=UpperCAmelCase      ,batch_size=UpperCAmelCase						)

   return train_dataloader, eval_dataloader, test_dataloader







def 						_A      ( UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     []
   # Download the dataset
   A__       =     load_dataset('glue'      ,'mrpc'						)
   # Create our splits
   A__       =     StratifiedKFold(n_splits=int(args.num_folds						)						)
   # Initialize accelerator
   A__       =     Accelerator(cpu=args.cpu      ,mixed_precision=args.mixed_precision						)
   # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
   A__       =     config['lr']
   A__       =     int(config['num_epochs']						)
   A__       =     int(config['seed']						)
   A__       =     int(config['batch_size']						)

   A__       =     evaluate.load('glue'      ,'mrpc'						)

   # If the batch size is too big we use gradient accumulation
   A__       =     1
   if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
      A__       =     batch_size // MAX_GPU_BATCH_SIZE
      A__       =     MAX_GPU_BATCH_SIZE

   set_seed(UpperCAmelCase						)

   # New Code #
   # Create our folds:
   A__       =     kfold.split(np.zeros(datasets['train'].num_rows						)      ,datasets['train']['label']						)
   A__       =     []
   # Iterate over them
   for i, (train_idxs, valid_idxs) in enumerate(UpperCAmelCase						):
      A__  ,			A__  ,			A__       =     get_fold_dataloaders(
          UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase      ,)
      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      A__       =     AutoModelForSequenceClassification.from_pretrained('bert-base-cased'      ,return_dict=UpperCAmelCase						)

      # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
      # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
      # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
      A__       =     model.to(accelerator.device						)

      # Instantiate optimizer
      A__       =     AdamW(params=model.parameters()      ,lr=UpperCAmelCase						)

      # Instantiate scheduler
      A__       =     get_linear_schedule_with_warmup(
          optimizer=UpperCAmelCase      ,num_warmup_steps=100      ,num_training_steps=(len(UpperCAmelCase						) * num_epochs) // gradient_accumulation_steps      ,)

      # 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.
      A__  ,			A__  ,			A__  ,			A__  ,			A__       =     accelerator.prepare(
          UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						)

      # Now we train the model
      for epoch in range(UpperCAmelCase						):
         model.train()
         for step, batch in enumerate(UpperCAmelCase						):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device						)
            A__       =     model(**UpperCAmelCase						)
            A__       =     outputs.loss
            A__       =     loss / gradient_accumulation_steps
            accelerator.backward(UpperCAmelCase						)
            if step % gradient_accumulation_steps == 0:
               optimizer.step()
               lr_scheduler.step()
               optimizer.zero_grad()

         model.eval()
         for step, batch in enumerate(UpperCAmelCase						):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device						)
            with torch.no_grad():
               A__       =     model(**UpperCAmelCase						)
            A__       =     outputs.logits.argmax(dim=-1						)
            A__  ,			A__       =     accelerator.gather_for_metrics((predictions, batch['labels'])						)
            metric.add_batch(
                predictions=UpperCAmelCase      ,references=UpperCAmelCase      ,)

         A__       =     metric.compute()
         # Use accelerator.print to print only on the main process.
         accelerator.print(F"""epoch {epoch}:"""      ,UpperCAmelCase						)

      # New Code #
      # We also run predictions on the test set at the very end
      A__       =     []
      for step, batch in enumerate(UpperCAmelCase						):
         # We could avoid this line since we set the accelerator with `device_placement=True`.
         batch.to(accelerator.device						)
         with torch.no_grad():
            A__       =     model(**UpperCAmelCase						)
         A__       =     outputs.logits
         A__  ,			A__       =     accelerator.gather_for_metrics((predictions, batch['labels'])						)
         fold_predictions.append(predictions.cpu()						)
         if i == 0:
            # We need all of the test predictions
            test_references.append(references.cpu()						)
        # Use accelerator.print to print only on the main process.
      test_predictions.append(torch.cat(UpperCAmelCase      ,dim=0						)						)
      # We now need to release all our memory and get rid of the current model, optimizer, etc
      accelerator.free_memory()
   # New Code #
   # Finally we check the accuracy of our folded results:
   A__       =     torch.cat(UpperCAmelCase      ,dim=0						)
   A__       =     torch.stack(UpperCAmelCase      ,dim=0						).sum(dim=0						).div(int(args.num_folds						)						).argmax(dim=-1						)
   A__       =     metric.compute(predictions=UpperCAmelCase      ,references=UpperCAmelCase						)
   accelerator.print('Average test metrics from all folds:'      ,UpperCAmelCase						)







def 						_A      ( ):






   '''simple docstring'''







   A__       =     argparse.ArgumentParser(description='Simple example of training script.'						)
   parser.add_argument(
       '--mixed_precision'      ,type=UpperCAmelCase      ,default=UpperCAmelCase      ,choices=['no', 'fp16', 'bf16', 'fp8']      ,help='Whether to use mixed precision. Choose'
       'between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'
       'and an Nvidia Ampere GPU.'      ,)
   parser.add_argument('--cpu'      ,action='store_true'      ,help='If passed, will train on the CPU.'						)
   # New Code #
   parser.add_argument('--num_folds'      ,type=UpperCAmelCase      ,default=3      ,help='The number of splits to perform across the dataset'						)
   A__       =     parser.parse_args()
   A__       =     {'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
   training_function(UpperCAmelCase      ,UpperCAmelCase						)


if __name__ == "__main__":
    main()


