


import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup


UpperCAmelCase     =       logging.getLogger(__name__)






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Union[str, Any]:




      """simple docstring"""

      random.seed(SCREAMING_SNAKE_CASE   )
      np.random.seed(SCREAMING_SNAKE_CASE   )
      torch.manual_seed(SCREAMING_SNAKE_CASE   )
      torch.cuda.manual_seed_all(SCREAMING_SNAKE_CASE   )






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )->    List[str]:




      """simple docstring"""

      model.eval()
      snake_case_			= next(model.parameters()   ).device
      snake_case_			= 1
      snake_case_			= torch.zeros((eval_batch_size, context_len)   ,   dtype=torch.long   ,   device=SCREAMING_SNAKE_CASE   )
      snake_case_			= DataLoader(SCREAMING_SNAKE_CASE   ,   shuffle=SCREAMING_SNAKE_CASE   ,   batch_size=SCREAMING_SNAKE_CASE   )
      snake_case_			= torch.zeros(1   ,   device=SCREAMING_SNAKE_CASE   )
      snake_case_			= 0
      for batch in eval_dataloader:
            batch.to(SCREAMING_SNAKE_CASE   )
            # pad
            context.zero_()
            for i in range(SCREAMING_SNAKE_CASE   ):
                  snake_case_			= batch[i]
            snake_case_			= model(SCREAMING_SNAKE_CASE   ,   labels=SCREAMING_SNAKE_CASE   )
            eval_loss += outputs[0].sum().item()
            nb_eval_examples += batch.size(0   )
      snake_case_			= eval_loss / nb_eval_examples
      snake_case_			= torch.exp(SCREAMING_SNAKE_CASE   )
      model.train()
      return perplexity






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE="gpt2"   )->    Optional[Any]:




      """simple docstring"""

      snake_case_			= GPTaLMHeadModel.from_pretrained(SCREAMING_SNAKE_CASE   ,   output_hidden_states=SCREAMING_SNAKE_CASE   )
      torch.save(model.state_dict()   ,   model_name + '''local.pt'''   )
      return model






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )->    Dict:




      """simple docstring"""

      snake_case_			= copy.deepcopy(SCREAMING_SNAKE_CASE   )
      model.to(SCREAMING_SNAKE_CASE   )

      snake_case_			= ['''bias''', '''LayerNorm.weight''']
      snake_case_			= [
          {
              '''params''': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay   )],
              '''weight_decay''': 0.0,
          },
          {'''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay   )], '''weight_decay''': 0.0},
      ]
      snake_case_			= AdamW(SCREAMING_SNAKE_CASE   ,   lr=5E-5   ,   eps=1E-8   )
      snake_case_			= get_linear_schedule_with_warmup(SCREAMING_SNAKE_CASE   ,   0   ,   SCREAMING_SNAKE_CASE   )
      torch.cuda.empty_cache()
      return model, lm_optimizer, lm_scheduler






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )->    int:




      """simple docstring"""

      snake_case_			= np.array(SCREAMING_SNAKE_CASE   ).mean()
      snake_case_			= np.array(SCREAMING_SNAKE_CASE   ).std()
      snake_case_			= (real_perps - avg) / std
      snake_case_			= list(zip(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   )
      joblib.dump(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE="dev.jbl"   ,   SCREAMING_SNAKE_CASE=recopy_gpta   ,   )->    int:




      """simple docstring"""

      snake_case_			= []
      snake_case_			= []
      snake_case_			= []

      # Initialize the transformer model
      snake_case_			= copy.deepcopy(SCREAMING_SNAKE_CASE   )
      orig_model.to(device='''cpu'''   )
      torch.cuda.empty_cache()

      # Compute perplexity of initial transformer model for comparison
      model.train()
      snake_case_			,       snake_case_			,       snake_case_			= recopy_model(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )

      for step in tqdm(range(SCREAMING_SNAKE_CASE   )   ):
            snake_case_			= torch.zeros((1, context_len)   ,   dtype=torch.long   ,   device=SCREAMING_SNAKE_CASE   )
            snake_case_			= random.choice(SCREAMING_SNAKE_CASE   )
            snake_case_			= random.randint(0   ,   len(story[0]   ) - context_len - 1   )
            snake_case_			= story[0][start : start + context_len]
            lm_optimizer.zero_grad()
            snake_case_			= model(SCREAMING_SNAKE_CASE   ,   labels=SCREAMING_SNAKE_CASE   )
            snake_case_			= outputs[0]
            snake_case_			= compute_perplexity(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
            model.train()
            lm_loss.backward()
            # Do LM backprop
            torch.nn.utils.clip_grad_norm_(model.parameters()   ,   3.0   )
            lm_optimizer.step()
            lm_scheduler.step()  # Update learning rate schedule

            # Compute perplexity after back-propagating on the selected context
            snake_case_			= compute_perplexity(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )

            # Periodically save the stored (X, IG(X)) pairs
            if step % 1000 == 0 and step > 1:
                  intermittent_save(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )

            # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
            snake_case_			,       snake_case_			,       snake_case_			= recopy_model(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )

            past_perps.append(past_perp.item()   )
            real_perps.append(orig_perp - real_perp.item()   )
            contexts.append(np.array(context.cpu()   )   )

      intermittent_save(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE="data/tokenized_stories_train_wikitext103.jbl"   ,   SCREAMING_SNAKE_CASE=100   ,   SCREAMING_SNAKE_CASE=1026   ,   SCREAMING_SNAKE_CASE=True   )->    Optional[int]:




      """simple docstring"""

      snake_case_			= joblib.load(SCREAMING_SNAKE_CASE   )
      print('''data loaded'''   )
      snake_case_			= []
      if trim:
            for i, example in enumerate(SCREAMING_SNAKE_CASE   ):
                  if len(example[0]   ) > min_len:
                        snake_case_			= random.randint(0   ,   len(example[0]   ) - context_len - 1   )
                        objective_set.append(example[0, start : start + context_len]   )
                  if len(SCREAMING_SNAKE_CASE   ) >= number:
                        break
            snake_case_			= []
            for j in range(i + 1   ,   len(SCREAMING_SNAKE_CASE   )   ):
                  if len(data[j][0]   ) > min_len:
                        train_data.append(data[j]   )
      else:
            snake_case_			= data[0:number]
            snake_case_			= data[number:]

      joblib.dump(SCREAMING_SNAKE_CASE   ,   '''objective_set.jbl'''   )
      print('''objective set saved'''   )
      return train_data, objective_set






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE=50   ,   SCREAMING_SNAKE_CASE="secondary_learner.pt"   )->    int:




      """simple docstring"""

      snake_case_			= torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu'''   )
      # We will use the first 512 pairs from our dataset as a test set for
      # our secondary learner and the rest to train
      snake_case_			= train_dataset[:512]
      snake_case_			= train_dataset[512:]
      snake_case_			= DataLoader(SCREAMING_SNAKE_CASE   ,   shuffle=SCREAMING_SNAKE_CASE   ,   batch_size=SCREAMING_SNAKE_CASE   )
      snake_case_			= DataLoader(SCREAMING_SNAKE_CASE   ,   shuffle=SCREAMING_SNAKE_CASE   ,   batch_size=SCREAMING_SNAKE_CASE   )

      # secondary learner model set up
      snake_case_			= nn.MSELoss()
      snake_case_			= nn.MSELoss(reduction='''sum'''   )
      secondary_learner.to(SCREAMING_SNAKE_CASE   )
      snake_case_			= torch.optim.Adam(secondary_learner.parameters()   ,   lr=0.0_0_0_0_1   )
      secondary_learner.train()

      # TODO in original code this is written as number of actual batches seen
      # not number of items seen but other places it is number of items instead.
      # improve consistency! changed this to epochs for clarity
      snake_case_			= float('''inf'''   )
      # Iterate through batches until we've used max_steps batches
      for epoch in range(int(SCREAMING_SNAKE_CASE   )   ):
            snake_case_			= 0.0
            secondary_learner.train()
            for step, batch in enumerate(SCREAMING_SNAKE_CASE   ):
                  snake_case_			= batch[0].to(SCREAMING_SNAKE_CASE   )
                  snake_case_			= batch[1].to(SCREAMING_SNAKE_CASE   )
                  snake_case_			= secondary_learner(SCREAMING_SNAKE_CASE   )
                  q_optimizer.zero_grad()
                  snake_case_			= loss(SCREAMING_SNAKE_CASE   ,   real_q.float()   )
                  q_loss.backward()
                  q_optimizer.step()
                  tr_q_loss += q_loss.item()

                  # model trains fairly quickly so we won't wait for a full epoch
                  # eval is triggered at eval_freq and end of epochs
                  if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(SCREAMING_SNAKE_CASE   )):
                        snake_case_			= tr_q_loss / (step + 1)

                        secondary_learner.eval()
                        snake_case_			= 0.0
                        snake_case_			= 0.0
                        snake_case_			= []
                        snake_case_			= []
                        # Compute performance of the secondary learner after this batch
                        for stepa, batcha in enumerate(SCREAMING_SNAKE_CASE   ):
                              snake_case_			= batcha[0].to(SCREAMING_SNAKE_CASE   )
                              snake_case_			= batcha[1].to(SCREAMING_SNAKE_CASE   )
                              snake_case_			= secondary_learner(SCREAMING_SNAKE_CASE   )
                              q_lossa += test_loss(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ).item()
                              sum_qa += torch.sum(SCREAMING_SNAKE_CASE   ).item()
                              for ei, i in enumerate(predicted_qa.cpu().detach().numpy()   ):
                                    predicted.append(i.item()   )
                              for ei, i in enumerate(real_qa.cpu().detach().numpy()   ):
                                    actual.append(i.item()   )

                        q_lossa /= len(SCREAMING_SNAKE_CASE   )
                        print(
                            '''Epoch: '''   ,   SCREAMING_SNAKE_CASE   ,   '''step: '''   ,   SCREAMING_SNAKE_CASE   ,   '''Avg. q:'''   ,   sum_qa / len(SCREAMING_SNAKE_CASE   )   ,   '''Train Loss: '''   ,   SCREAMING_SNAKE_CASE   ,   '''Test Loss: '''   ,   SCREAMING_SNAKE_CASE   ,   )
                        if q_lossa < best_test_loss:
                              joblib.dump((predicted, actual)   ,   '''pred_vs_actual.jbl'''   )
                              torch.save(secondary_learner.state_dict()   ,   SCREAMING_SNAKE_CASE   )
                              snake_case_			= q_lossa

                  secondary_learner.train()
      return secondary_learner







class lowerCAmelCase_		(					nn.Module			):

 '''simple docstring'''

 def __init__(      self						, _UpperCAmelCase   ):
       super(_UpperCAmelCase						, self   ).__init__()
       snake_case_			= model.transformer.wte
       snake_case_			= copy.deepcopy(model.transformer.wte.weight   )
       snake_case_			= nn.Convad(self.embeddings.weight.size(1   )						, 2_56						, 3						, padding=1   )
       snake_case_			= nn.Sequential(nn.Linear(2_56						, 32   )						, nn.Dropout(p=0.1   )						, nn.Linear(32						, 32   )						, nn.Linear(32						, 1   )   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase   ):
       snake_case_			= torch.max(self.conv(self.embeddings(_UpperCAmelCase   ).squeeze(1   ).transpose(1						, 2   )   )						, 2   )[0]
       snake_case_			= self.fc(_UpperCAmelCase   )
       return qs.squeeze(1   )





 @classmethod
 def 		UpperCamelCase__	(      cls						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= cls(_UpperCAmelCase   )  # this calls __init__
       snake_case_			= torch.load(_UpperCAmelCase   )
       secondary_learner.load_state_dict(_UpperCAmelCase   )
       snake_case_			= model.transformer.wte
       snake_case_			= copy.deepcopy(model.transformer.wte.weight   )
       return secondary_learner



import argparse
import json
from pathlib import Path

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

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase     =       logging.get_logger(__name__)






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Tuple:




      """simple docstring"""

      snake_case_			= ConvNextConfig()

      if "tiny" in checkpoint_url:
            snake_case_			= [3, 3, 9, 3]
            snake_case_			= [96, 192, 384, 768]
      if "small" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [96, 192, 384, 768]
      if "base" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [128, 256, 512, 1024]
      if "large" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [192, 384, 768, 1536]
      if "xlarge" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [256, 512, 1024, 2048]

      if "1k" in checkpoint_url:
            snake_case_			= 1000
            snake_case_			= '''imagenet-1k-id2label.json'''
            snake_case_			= (1, 1000)
      else:
            snake_case_			= 2_1841
            snake_case_			= '''imagenet-22k-id2label.json'''
            snake_case_			= (1, 2_1841)

      snake_case_			= '''huggingface/label-files'''
      snake_case_			= num_labels
      snake_case_			= json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   repo_type='''dataset'''   )   ,   '''r'''   )   )
      snake_case_			= {int(SCREAMING_SNAKE_CASE   ): v for k, v in idalabel.items()}
      if "1k" not in checkpoint_url:
            # this dataset contains 21843 labels but the model only has 21841
            # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
            del idalabel[9205]
            del idalabel[1_5027]
      snake_case_			= idalabel
      snake_case_			= {v: k for k, v in idalabel.items()}
      snake_case_			= hidden_sizes
      snake_case_			= depths

      return config, expected_shape






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Tuple:




      """simple docstring"""

      if "downsample_layers.0.0" in name:
            snake_case_			= name.replace('''downsample_layers.0.0'''   ,   '''embeddings.patch_embeddings'''   )
      if "downsample_layers.0.1" in name:
            snake_case_			= name.replace('''downsample_layers.0.1'''   ,   '''embeddings.norm'''   )  # we rename to layernorm later on
      if "downsample_layers.1.0" in name:
            snake_case_			= name.replace('''downsample_layers.1.0'''   ,   '''stages.1.downsampling_layer.0'''   )
      if "downsample_layers.1.1" in name:
            snake_case_			= name.replace('''downsample_layers.1.1'''   ,   '''stages.1.downsampling_layer.1'''   )
      if "downsample_layers.2.0" in name:
            snake_case_			= name.replace('''downsample_layers.2.0'''   ,   '''stages.2.downsampling_layer.0'''   )
      if "downsample_layers.2.1" in name:
            snake_case_			= name.replace('''downsample_layers.2.1'''   ,   '''stages.2.downsampling_layer.1'''   )
      if "downsample_layers.3.0" in name:
            snake_case_			= name.replace('''downsample_layers.3.0'''   ,   '''stages.3.downsampling_layer.0'''   )
      if "downsample_layers.3.1" in name:
            snake_case_			= name.replace('''downsample_layers.3.1'''   ,   '''stages.3.downsampling_layer.1'''   )
      if "stages" in name and "downsampling_layer" not in name:
            # stages.0.0. for instance should be renamed to stages.0.layers.0.
            snake_case_			= name[: len('''stages.0'''   )] + '''.layers''' + name[len('''stages.0'''   ) :]
      if "stages" in name:
            snake_case_			= name.replace('''stages'''   ,   '''encoder.stages'''   )
      if "norm" in name:
            snake_case_			= name.replace('''norm'''   ,   '''layernorm'''   )
      if "gamma" in name:
            snake_case_			= name.replace('''gamma'''   ,   '''layer_scale_parameter'''   )
      if "head" in name:
            snake_case_			= name.replace('''head'''   ,   '''classifier'''   )

      return name






def 			__lowerCAmelCase       ()->    Optional[Any]:




      """simple docstring"""

      snake_case_			= '''http://images.cocodataset.org/val2017/000000039769.jpg'''
      snake_case_			= Image.open(requests.get(SCREAMING_SNAKE_CASE   ,   stream=SCREAMING_SNAKE_CASE   ).raw   )
      return im






@torch.no_grad()
def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )->    str:




      """simple docstring"""

      snake_case_			,       snake_case_			= get_convnext_config(SCREAMING_SNAKE_CASE   )
      # load original state_dict from URL
      snake_case_			= torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE   )['''model''']
      # rename keys
      for key in state_dict.copy().keys():
            snake_case_			= state_dict.pop(SCREAMING_SNAKE_CASE   )
            snake_case_			= val
      # add prefix to all keys expect classifier head
      for key in state_dict.copy().keys():
            snake_case_			= state_dict.pop(SCREAMING_SNAKE_CASE   )
            if not key.startswith('''classifier'''   ):
                  snake_case_			= '''convnext.''' + key
            snake_case_			= val

      # load HuggingFace model
      snake_case_			= ConvNextForImageClassification(SCREAMING_SNAKE_CASE   )
      model.load_state_dict(SCREAMING_SNAKE_CASE   )
      model.eval()

      # Check outputs on an image, prepared by ConvNextImageProcessor
      snake_case_			= 224 if '''224''' in checkpoint_url else 384
      snake_case_			= ConvNextImageProcessor(size=SCREAMING_SNAKE_CASE   )
      snake_case_			= image_processor(images=prepare_img()   ,   return_tensors='''pt'''   ).pixel_values

      snake_case_			= model(SCREAMING_SNAKE_CASE   ).logits

      # note: the logits below were obtained without center cropping
      if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.1_2_1_0, -0.6_6_0_5, 0.1_9_1_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.4_4_7_3, -0.1_8_4_7, -0.6_3_6_5]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
            snake_case_			= torch.tensor([0.4_5_2_5, 0.7_5_3_9, 0.0_3_0_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
            snake_case_			= torch.tensor([0.3_5_6_1, 0.6_3_5_0, -0.0_3_8_4]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
            snake_case_			= torch.tensor([0.4_1_7_4, -0.0_9_8_9, 0.1_4_8_9]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
            snake_case_			= torch.tensor([0.2_5_1_3, -0.1_3_4_9, -0.1_6_1_3]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
            snake_case_			= torch.tensor([1.2_9_8_0, 0.3_6_3_1, -0.1_1_9_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
            snake_case_			= torch.tensor([1.2_9_6_3, 0.1_2_2_7, 0.1_7_2_3]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
            snake_case_			= torch.tensor([1.7_9_5_6, 0.8_3_9_0, 0.2_8_2_0]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
            snake_case_			= torch.tensor([-0.2_8_2_2, -0.0_5_0_2, -0.0_8_7_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
            snake_case_			= torch.tensor([-0.5_6_7_2, -0.0_7_3_0, -0.4_3_4_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
            snake_case_			= torch.tensor([0.2_6_8_1, 0.2_3_6_5, 0.6_2_4_6]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
            snake_case_			= torch.tensor([-0.2_6_4_2, 0.3_9_3_1, 0.5_1_1_6]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.6_6_7_7, -0.1_8_7_3, -0.8_3_7_9]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
            snake_case_			= torch.tensor([-0.7_7_4_9, -0.2_9_6_7, -0.6_4_4_4]   )
      else:
            raise ValueError(f'''Unknown URL: {checkpoint_url}'''   )

      assert torch.allclose(logits[0, :3]   ,   SCREAMING_SNAKE_CASE   ,   atol=1E-3   )
      assert logits.shape == expected_shape

      Path(SCREAMING_SNAKE_CASE   ).mkdir(exist_ok=SCREAMING_SNAKE_CASE   )
      print(f'''Saving model to {pytorch_dump_folder_path}'''   )
      model.save_pretrained(SCREAMING_SNAKE_CASE   )
      print(f'''Saving image processor to {pytorch_dump_folder_path}'''   )
      image_processor.save_pretrained(SCREAMING_SNAKE_CASE   )

      print('''Pushing model to the hub...'''   )
      snake_case_			= '''convnext'''
      if "tiny" in checkpoint_url:
            model_name += "-tiny"
      elif "small" in checkpoint_url:
            model_name += "-small"
      elif "base" in checkpoint_url:
            model_name += "-base"
      elif "xlarge" in checkpoint_url:
            model_name += "-xlarge"
      elif "large" in checkpoint_url:
            model_name += "-large"
      if "224" in checkpoint_url:
            model_name += "-224"
      elif "384" in checkpoint_url:
            model_name += "-384"
      if "22k" in checkpoint_url and "1k" not in checkpoint_url:
            model_name += "-22k"
      if "22k" in checkpoint_url and "1k" in checkpoint_url:
            model_name += "-22k-1k"

      model.push_to_hub(
          repo_path_or_name=Path(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   ,   organization='''nielsr'''   ,   commit_message='''Add model'''   ,   )


if __name__ == "__main__":
   UpperCAmelCase     =       argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--checkpoint_url""",
       default="""https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth""",
       type=str,
       help="""URL of the original ConvNeXT checkpoint you'd like to convert.""",
   )
   parser.add_argument(
       """--pytorch_dump_folder_path""",
       default=None,
       type=str,
       required=True,
       help="""Path to the output PyTorch model directory.""",
   )

   UpperCAmelCase     =       parser.parse_args()
   convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)