

import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__	):
						snake_case_       =     ConvNextVaConfig()

						if "atto" in checkpoint_url:
												snake_case_       =     [2, 2, 6, 2]
												snake_case_       =     [40, 80, 160, 320]
						if "femto" in checkpoint_url:
												snake_case_       =     [2, 2, 6, 2]
												snake_case_       =     [48, 96, 192, 384]
						if "pico" in checkpoint_url:
												snake_case_       =     [2, 2, 6, 2]
												snake_case_       =     [64, 128, 256, 512]
						if "nano" in checkpoint_url:
												snake_case_       =     [2, 2, 8, 2]
												snake_case_       =     [80, 160, 320, 640]
						if "tiny" in checkpoint_url:
												snake_case_       =     [3, 3, 9, 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 "huge" in checkpoint_url:
												snake_case_       =     [3, 3, 27, 3]
												snake_case_       =     [352, 704, 1408, 2816]

						snake_case_       =     1000
						snake_case_       =     '''imagenet-1k-id2label.json'''
						snake_case_       =     (1, 1000)

						snake_case_       =     '''huggingface/label-files'''
						snake_case_       =     num_labels
						snake_case_       =     json.load(open(hf_hub_download(UpperCamelCase_  ,		UpperCamelCase_  ,		repo_type='''dataset'''	)  ,		'''r'''	)	)
						snake_case_       =     {int(UpperCamelCase_	): v for k, v in idalabel.items()}

						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   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__	):



						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 "gamma" in name:
												snake_case_       =     name.replace('''gamma'''  ,		'''weight'''	)
						if "beta" in name:
												snake_case_       =     name.replace('''beta'''  ,		'''bias'''	)
						if "stages" in name:
												snake_case_       =     name.replace('''stages'''  ,		'''encoder.stages'''	)
						if "norm" in name:
												snake_case_       =     name.replace('''norm'''  ,		'''layernorm'''	)
						if "head" in name:
												snake_case_       =     name.replace('''head'''  ,		'''classifier'''	)

						return name


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


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__	):



						if "224" in checkpoint_url:
												snake_case_       =     224
												snake_case_       =     224 / 256
						elif "384" in checkpoint_url:
												snake_case_       =     384
												snake_case_       =     None
						else:
												snake_case_       =     512
												snake_case_       =     None

						return ConvNextImageProcessor(
						    size=UpperCamelCase_  ,		crop_pct=UpperCamelCase_  ,		image_mean=[0.485, 0.456, 0.406]  ,		image_std=[0.229, 0.224, 0.225]  ,		resample=PILImageResampling.BICUBIC  ,		)


@torch.no_grad()
def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
						print('''Downloading original model from checkpoint...'''	)
						# define ConvNeXTV2 configuration based on URL
						snake_case_,   snake_case_       =     get_convnextva_config(UpperCamelCase_	)
						# load original state_dict from URL
						snake_case_       =     torch.hub.load_state_dict_from_url(UpperCamelCase_	)['''model''']

						print('''Converting model parameters...'''	)
						# rename keys
						for key in state_dict.copy().keys():
												snake_case_       =     state_dict.pop(UpperCamelCase_	)
												snake_case_       =     val
						# add prefix to all keys expect classifier head
						for key in state_dict.copy().keys():
												snake_case_       =     state_dict.pop(UpperCamelCase_	)
												if not key.startswith('''classifier'''	):
																		snake_case_       =     '''convnextv2.''' + key
												snake_case_       =     val

						# load HuggingFace model
						snake_case_       =     ConvNextVaForImageClassification(UpperCamelCase_	)
						model.load_state_dict(UpperCamelCase_	)
						model.eval()

						# Check outputs on an image, prepared by ConvNextImageProcessor
						snake_case_       =     convert_preprocessor(UpperCamelCase_	)
						snake_case_       =     preprocessor(images=prepare_img()  ,		return_tensors='''pt'''	)
						snake_case_       =     model(**UpperCamelCase_	).logits

						# note: the logits below were obtained without center cropping
						if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
												snake_case_       =     torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
												snake_case_       =     torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
												snake_case_       =     torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
												snake_case_       =     torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
												snake_case_       =     torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
												snake_case_       =     torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
												snake_case_       =     torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
												snake_case_       =     torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
												snake_case_       =     torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
												snake_case_       =     torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]	)
						else:
												raise ValueError(F'''Unknown URL: {checkpoint_url}'''	)

						assert torch.allclose(logits[0, :5]  ,		UpperCamelCase_  ,		atol=1E-3	)
						assert logits.shape == expected_shape
						print('''Model outputs match the original results!'''	)

						if save_model:
												print('''Saving model to local...'''	)
												# Create folder to save model
												if not os.path.isdir(UpperCamelCase_	):
																		os.mkdir(UpperCamelCase_	)

												model.save_pretrained(UpperCamelCase_	)
												preprocessor.save_pretrained(UpperCamelCase_	)

						snake_case_       =     '''convnextv2'''
						if "atto" in checkpoint_url:
												model_name += "-atto"
						if "femto" in checkpoint_url:
												model_name += "-femto"
						if "pico" in checkpoint_url:
												model_name += "-pico"
						if "nano" in checkpoint_url:
												model_name += "-nano"
						elif "tiny" in checkpoint_url:
												model_name += "-tiny"
						elif "base" in checkpoint_url:
												model_name += "-base"
						elif "large" in checkpoint_url:
												model_name += "-large"
						elif "huge" in checkpoint_url:
												model_name += "-huge"
						if "22k" in checkpoint_url and "1k" not in checkpoint_url:
												model_name += "-22k"
						elif "22k" in checkpoint_url and "1k" in checkpoint_url:
												model_name += "-22k-1k"
						elif "1k" in checkpoint_url:
												model_name += "-1k"
						if "224" in checkpoint_url:
												model_name += "-224"
						elif "384" in checkpoint_url:
												model_name += "-384"
						elif "512" in checkpoint_url:
												model_name += "-512"

						if push_to_hub:
												print(F'''Pushing {model_name} to the hub...'''	)
												model.push_to_hub(UpperCamelCase_	)
												preprocessor.push_to_hub(UpperCamelCase_	)


if __name__ == "__main__":
					lowerCAmelCase_	     =      argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--checkpoint_url''',
					    default='''https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt''',
					    type=str,
					    help='''URL of the original ConvNeXTV2 checkpoint you\'d like to convert.''',
					)
					parser.add_argument(
					    '''--pytorch_dump_folder_path''',
					    default='''model''',
					    type=str,
					    help='''Path to the output PyTorch model directory.''',
					)
					parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
					parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image preprocessor to the hub''')

					lowerCAmelCase_	     =      parser.parse_args()
					convert_convnextva_checkpoint(
					    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
					)


import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_splinter import SplinterTokenizer


lowerCAmelCase_	     =     logging.get_logger(__name__)

lowerCAmelCase_	     =     {'''vocab_file''': '''vocab.txt'''}

lowerCAmelCase_	     =     {
    '''vocab_file''': {
        '''tau/splinter-base''': '''https://huggingface.co/tau/splinter-base/resolve/main/vocab.txt''',
        '''tau/splinter-base-qass''': '''https://huggingface.co/tau/splinter-base-qass/resolve/main/vocab.txt''',
        '''tau/splinter-large''': '''https://huggingface.co/tau/splinter-large/resolve/main/vocab.txt''',
        '''tau/splinter-large-qass''': '''https://huggingface.co/tau/splinter-large-qass/resolve/main/vocab.txt''',
    }
}

lowerCAmelCase_	     =     {
    '''tau/splinter-base''': 5_12,
    '''tau/splinter-base-qass''': 5_12,
    '''tau/splinter-large''': 5_12,
    '''tau/splinter-large-qass''': 5_12,
}

lowerCAmelCase_	     =     {
    '''tau/splinter-base''': {'''do_lower_case''': False},
    '''tau/splinter-base-qass''': {'''do_lower_case''': False},
    '''tau/splinter-large''': {'''do_lower_case''': False},
    '''tau/splinter-large-qass''': {'''do_lower_case''': False},
}





class      snake_case_       (		__A ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE					:       List[str]						      =    VOCAB_FILES_NAMES
			SCREAMING_SNAKE_CASE					:       List[str]						      =    PRETRAINED_VOCAB_FILES_MAP
			SCREAMING_SNAKE_CASE					:       Optional[int]						      =    PRETRAINED_INIT_CONFIGURATION
			SCREAMING_SNAKE_CASE					:       List[str]						      =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			SCREAMING_SNAKE_CASE					:       Any						      =    SplinterTokenizer







			def __init__(					self			:    List[str]							,       _UpperCamelCase			:    Dict=None							,       _UpperCamelCase			:    Union[str, Any]=None							,       _UpperCamelCase			:    Union[str, Any]=True							,       _UpperCamelCase			:    List[Any]="[UNK]"							,       _UpperCamelCase			:    List[str]="[SEP]"							,       _UpperCamelCase			:    Optional[int]="[PAD]"							,       _UpperCamelCase			:    Tuple="[CLS]"							,       _UpperCamelCase			:    Optional[int]="[MASK]"							,       _UpperCamelCase			:    str="[QUESTION]"							,       _UpperCamelCase			:    Optional[Any]=True							,       _UpperCamelCase			:    Dict=None							,       **_UpperCamelCase			:    Optional[int]							,       )				->Union[str, Any]:
									super().__init__(
									    _UpperCamelCase							,       tokenizer_file=_UpperCamelCase							,       do_lower_case=_UpperCamelCase							,       unk_token=_UpperCamelCase							,       sep_token=_UpperCamelCase							,       pad_token=_UpperCamelCase							,       cls_token=_UpperCamelCase							,       mask_token=_UpperCamelCase							,       tokenize_chinese_chars=_UpperCamelCase							,       strip_accents=_UpperCamelCase							,       additional_special_tokens=(question_token,)							,       **_UpperCamelCase							,       )

									snake_case_       =     json.loads(self.backend_tokenizer.normalizer.__getstate__()							)
									if (
									    pre_tok_state.get('''lowercase'''							,       _UpperCamelCase							) != do_lower_case
									    or pre_tok_state.get('''strip_accents'''							,       _UpperCamelCase							) != strip_accents
									):
															snake_case_       =     getattr(_UpperCamelCase							,       pre_tok_state.pop('''type'''							)							)
															snake_case_       =     do_lower_case
															snake_case_       =     strip_accents
															snake_case_       =     pre_tok_class(**_UpperCamelCase							)

									snake_case_       =     do_lower_case







			@property
			def     snake_case__(					self			:    Optional[int]							)				->int:
									return self.convert_tokens_to_ids(self.question_token							)







			def     snake_case__(					self			:    Dict							,       _UpperCamelCase			:    List[int]							,       _UpperCamelCase			:    Optional[List[int]] = None							)				->List[int]:
									if token_ids_a is None:
															return [self.cls_token_id] + token_ids_a + [self.sep_token_id]

									snake_case_       =     [self.cls_token_id]
									snake_case_       =     [self.sep_token_id]
									snake_case_       =     [self.question_token_id] + [self.convert_tokens_to_ids('''.'''							)]
									if self.padding_side == "right":
															# Input is question-then-context
															return cls + token_ids_a + question_suffix + sep + token_ids_a + sep
									else:
															# Input is context-then-question
															return cls + token_ids_a + sep + token_ids_a + question_suffix + sep







			def     snake_case__(					self			:    str							,       _UpperCamelCase			:    List[int]							,       _UpperCamelCase			:    Optional[List[int]] = None							)				->List[int]:
									snake_case_       =     [self.sep_token_id]
									snake_case_       =     [self.cls_token_id]
									snake_case_       =     [self.question_token_id] + [self.convert_tokens_to_ids('''.'''							)]
									if token_ids_a is None:
															return len(cls + token_ids_a + sep							) * [0]

									if self.padding_side == "right":
															# Input is question-then-context
															return len(cls + token_ids_a + question_suffix + sep							) * [0] + len(token_ids_a + sep							) * [1]
									else:
															# Input is context-then-question
															return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + question_suffix + sep							) * [1]







			def     snake_case__(					self			:    Union[str, Any]							,       _UpperCamelCase			:    str							,       _UpperCamelCase			:    Optional[str] = None							)				->Tuple[str]:
									snake_case_       =     self._tokenizer.model.save(_UpperCamelCase							,       name=_UpperCamelCase							)
									return tuple(_UpperCamelCase							)