'''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 __lowerCAmelCase (							UpperCamelCase__				)	->    Optional[Any]:
    __lowerCamelCase		=  [2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
    __lowerCamelCase		=  True if '''large''' in model_name or '''huge''' in model_name else False
    __lowerCamelCase		=  True if '''large''' in model_name or '''huge''' in model_name else False
    __lowerCamelCase		=  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		=  [3, 3, 3, 3]
            __lowerCamelCase		=  [5, 5, 5, 5]
        elif "fl4" in model_name:
            __lowerCamelCase		=  [4, 4, 4, 4]
            __lowerCamelCase		=  [3, 3, 3, 3]

    if "tiny" in model_name or "small" in model_name or "base" in model_name:
        __lowerCamelCase		=  [3, 3, 3, 3]
        if "lrf" in model_name:
            __lowerCamelCase		=  [3, 3, 3, 3]
        else:
            __lowerCamelCase		=  [2, 2, 2, 2]

    if "tiny" in model_name:
        __lowerCamelCase		=  96
    elif "small" in model_name:
        __lowerCamelCase		=  96
    elif "base" in model_name:
        __lowerCamelCase		=  1_28
    elif "large" in model_name:
        __lowerCamelCase		=  1_92
    elif "xlarge" in model_name:
        __lowerCamelCase		=  2_56
    elif "huge" in model_name:
        __lowerCamelCase		=  3_52

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

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

    __lowerCamelCase		=  FocalNetConfig(
        embed_dim=UpperCamelCase__				,  depths=UpperCamelCase__				,  focal_levels=UpperCamelCase__				,  focal_windows=UpperCamelCase__				,  use_conv_embed=UpperCamelCase__				,  idalabel=UpperCamelCase__				,  labelaid=UpperCamelCase__				,  use_post_layernorm=UpperCamelCase__				,  use_layerscale=UpperCamelCase__				,  )

    return config






def __lowerCAmelCase (							UpperCamelCase__				)	->    Optional[Any]:
    if "patch_embed.proj" in name:
        __lowerCamelCase		=  name.replace('''patch_embed.proj'''				,  '''embeddings.patch_embeddings.projection'''				)
    if "patch_embed.norm" in name:
        __lowerCamelCase		=  name.replace('''patch_embed.norm'''				,  '''embeddings.norm'''				)
    if "layers" in name:
        __lowerCamelCase		=  '''encoder.''' + name
    if "encoder.layers" in name:
        __lowerCamelCase		=  name.replace('''encoder.layers'''				,  '''encoder.stages'''				)
    if "downsample.proj" in name:
        __lowerCamelCase		=  name.replace('''downsample.proj'''				,  '''downsample.projection'''				)
    if "blocks" in name:
        __lowerCamelCase		=  name.replace('''blocks'''				,  '''layers'''				)
    if "modulation.f.weight" in name or "modulation.f.bias" in name:
        __lowerCamelCase		=  name.replace('''modulation.f'''				,  '''modulation.projection_in'''				)
    if "modulation.h.weight" in name or "modulation.h.bias" in name:
        __lowerCamelCase		=  name.replace('''modulation.h'''				,  '''modulation.projection_context'''				)
    if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
        __lowerCamelCase		=  name.replace('''modulation.proj'''				,  '''modulation.projection_out'''				)

    if name == "norm.weight":
        __lowerCamelCase		=  '''layernorm.weight'''
    if name == "norm.bias":
        __lowerCamelCase		=  '''layernorm.bias'''

    if "head" in name:
        __lowerCamelCase		=  name.replace('''head'''				,  '''classifier'''				)
    else:
        __lowerCamelCase		=  '''focalnet.''' + name

    return name






def __lowerCAmelCase (							UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__=False				)	->    Optional[Any]:
    # fmt: off
    __lowerCamelCase		=  {
        '''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		=  model_name_to_url[model_name]
    print('''Checkpoint URL: '''				,  UpperCamelCase__				)
    __lowerCamelCase		=  torch.hub.load_state_dict_from_url(UpperCamelCase__				,  map_location='''cpu'''				)['''model''']

    # rename keys
    for key in state_dict.copy().keys():
        __lowerCamelCase		=  state_dict.pop(UpperCamelCase__				)
        __lowerCamelCase		=  val

    __lowerCamelCase		=  get_focalnet_config(UpperCamelCase__				)
    __lowerCamelCase		=  FocalNetForImageClassification(UpperCamelCase__				)
    model.eval()

    # load state dict
    model.load_state_dict(UpperCamelCase__				)

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

    __lowerCamelCase		=  BitImageProcessor(
        do_resize=UpperCamelCase__				,  size={'''shortest_edge''': 2_56}				,  resample=PILImageResampling.BILINEAR				,  do_center_crop=UpperCamelCase__				,  crop_size=2_24				,  do_normalize=UpperCamelCase__				,  image_mean=UpperCamelCase__				,  image_std=UpperCamelCase__				,  )
    __lowerCamelCase		=  Image.open(requests.get(UpperCamelCase__				,  stream=UpperCamelCase__				).raw				)
    __lowerCamelCase		=  processor(images=UpperCamelCase__				,  return_tensors='''pt'''				)

    __lowerCamelCase		=  transforms.Compose(
        [
            transforms.Resize(2_56				),
            transforms.CenterCrop(2_24				),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]				,  std=[0.2_2_9, 0.2_2_4, 0.2_2_5]				),
        ]				)

    __lowerCamelCase		=  image_transforms(UpperCamelCase__				).unsqueeze(0				)

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

    __lowerCamelCase		=  model(**UpperCamelCase__				)

    __lowerCamelCase		=  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		=  torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]				)
    elif model_name == "focalnet-tiny-lrf":
        __lowerCamelCase		=  torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]				)
    elif model_name == "focalnet-small":
        __lowerCamelCase		=  torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]				)
    elif model_name == "focalnet-small-lrf":
        __lowerCamelCase		=  torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]				)
    elif model_name == "focalnet-base":
        __lowerCamelCase		=  torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]				)
    elif model_name == "focalnet-base-lrf":
        __lowerCamelCase		=  torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]				)
    assert torch.allclose(outputs.logits[0, :3]				,  UpperCamelCase__				,  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(UpperCamelCase__				)
        processor.save_pretrained(UpperCamelCase__				)

    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__":
  __UpperCAmelCase     							=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.",
  )

  __UpperCAmelCase     							=parser.parse_args()
  convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





'''simple docstring'''





import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_lowerCAmelCase      						=	logging.get_logger(__name__)
_lowerCAmelCase      						=	{'''vocab_file''': '''spiece.model'''}

_lowerCAmelCase      						=	{
    '''vocab_file''': {
        '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/spiece.model''',
        '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/spiece.model''',
        '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model''',
        '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model''',
        '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/spiece.model''',
        '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/spiece.model''',
        '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model''',
        '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model''',
    }
}

_lowerCAmelCase      						=	{
    '''albert-base-v1''': 512,
    '''albert-large-v1''': 512,
    '''albert-xlarge-v1''': 512,
    '''albert-xxlarge-v1''': 512,
    '''albert-base-v2''': 512,
    '''albert-large-v2''': 512,
    '''albert-xlarge-v2''': 512,
    '''albert-xxlarge-v2''': 512,
}

_lowerCAmelCase      						=	'''▁'''






class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_     ):
  '''simple docstring'''




  __lowercase							:      Any				=      VOCAB_FILES_NAMES
  __lowercase							:      Optional[Any]				=      PRETRAINED_VOCAB_FILES_MAP
  __lowercase							:      int				=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




  def __init__(    self  ,__UpperCAmelCase  ,__UpperCAmelCase=True  ,__UpperCAmelCase=True  ,__UpperCAmelCase=False  ,__UpperCAmelCase="[CLS]"  ,__UpperCAmelCase="[SEP]"  ,__UpperCAmelCase="<unk>"  ,__UpperCAmelCase="[SEP]"  ,__UpperCAmelCase="<pad>"  ,__UpperCAmelCase="[CLS]"  ,__UpperCAmelCase="[MASK]"  ,__UpperCAmelCase = None  ,**__UpperCAmelCase  ,)    ->		None:
      # Mask token behave like a normal word, i.e. include the space before it and
      # is included in the raw text, there should be a match in a non-normalized sentence.
      lowerCAmelCase__ :  Tuple							      =					(
          AddedToken(__UpperCAmelCase  ,lstrip=__UpperCAmelCase  ,rstrip=__UpperCAmelCase  ,normalized=__UpperCAmelCase      )
          if isinstance(__UpperCAmelCase  ,__UpperCAmelCase      )
          else mask_token
      )

      lowerCAmelCase__ :  Any							      =					{} if sp_model_kwargs is None else sp_model_kwargs

      super().__init__(
          do_lower_case=__UpperCAmelCase  ,remove_space=__UpperCAmelCase  ,keep_accents=__UpperCAmelCase  ,bos_token=__UpperCAmelCase  ,eos_token=__UpperCAmelCase  ,unk_token=__UpperCAmelCase  ,sep_token=__UpperCAmelCase  ,pad_token=__UpperCAmelCase  ,cls_token=__UpperCAmelCase  ,mask_token=__UpperCAmelCase  ,sp_model_kwargs=self.sp_model_kwargs  ,**__UpperCAmelCase  ,)

      lowerCAmelCase__ :  str							      =					do_lower_case
      lowerCAmelCase__ :  int							      =					remove_space
      lowerCAmelCase__ :  Tuple							      =					keep_accents
      lowerCAmelCase__ :  Any							      =					vocab_file

      lowerCAmelCase__ :  Tuple							      =					spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(__UpperCAmelCase      )




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Optional[int]:
      return len(self.sp_model      )




  def 							UpperCAmelCase_	(    self      )    ->		Optional[Any]:
      lowerCAmelCase__ :  Union[str, Any]							      =					{self.convert_ids_to_tokens(__UpperCAmelCase      ): i for i in range(self.vocab_size      )}
      vocab.update(self.added_tokens_encoder      )
      return vocab




  def __getstate__(    self      )    ->		Any:
      lowerCAmelCase__ :  Optional[Any]							      =					self.__dict__.copy()
      lowerCAmelCase__ :  Optional[Any]							      =					None
      return state




  def __setstate__(    self  ,__UpperCAmelCase      )    ->		List[Any]:
      lowerCAmelCase__ :  List[str]							      =					d

      # for backward compatibility
      if not hasattr(self  ,"""sp_model_kwargs"""      ):
          lowerCAmelCase__ :  Union[str, Any]							      =					{}

      lowerCAmelCase__ :  Any							      =					spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(self.vocab_file      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Optional[int]:
      if self.remove_space:
          lowerCAmelCase__ :  int							      =					""" """.join(inputs.strip().split()      )
      else:
          lowerCAmelCase__ :  str							      =					inputs
      lowerCAmelCase__ :  Tuple							      =					outputs.replace("""``"""  ,"""\""""      ).replace("""''"""  ,"""\""""      )

      if not self.keep_accents:
          lowerCAmelCase__ :  Any							      =					unicodedata.normalize("""NFKD"""  ,__UpperCAmelCase      )
          lowerCAmelCase__ :  Dict							      =					"""""".join([c for c in outputs if not unicodedata.combining(__UpperCAmelCase      )]      )
      if self.do_lower_case:
          lowerCAmelCase__ :  Tuple							      =					outputs.lower()

      return outputs




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		List[str]:
      lowerCAmelCase__ :  List[str]							      =					self.preprocess_text(__UpperCAmelCase      )
      lowerCAmelCase__ :  Optional[int]							      =					self.sp_model.encode(__UpperCAmelCase  ,out_type=__UpperCAmelCase      )
      lowerCAmelCase__ :  str							      =					[]
      for piece in pieces:
          if len(__UpperCAmelCase      ) > 1 and piece[-1] == str(""","""      ) and piece[-2].isdigit():
              lowerCAmelCase__ :  List[Any]							      =					self.sp_model.EncodeAsPieces(piece[:-1].replace(__UpperCAmelCase  ,""""""      )      )
              if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
                  if len(cur_pieces[0]      ) == 1:
                      lowerCAmelCase__ :  str							      =					cur_pieces[1:]
                  else:
                      lowerCAmelCase__ :  int							      =					cur_pieces[0][1:]
              cur_pieces.append(piece[-1]      )
              new_pieces.extend(__UpperCAmelCase      )
          else:
              new_pieces.append(__UpperCAmelCase      )

      return new_pieces




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		List[str]:
      return self.sp_model.PieceToId(__UpperCAmelCase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Any:
      return self.sp_model.IdToPiece(__UpperCAmelCase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		str:
      lowerCAmelCase__ :  str							      =					[]
      lowerCAmelCase__ :  Tuple							      =					""""""
      lowerCAmelCase__ :  Tuple							      =					False
      for token in tokens:
          # make sure that special tokens are not decoded using sentencepiece model
          if token in self.all_special_tokens:
              if not prev_is_special:
                  out_string += " "
              out_string += self.sp_model.decode(__UpperCAmelCase      ) + token
              lowerCAmelCase__ :  Union[str, Any]							      =					True
              lowerCAmelCase__ :  List[str]							      =					[]
          else:
              current_sub_tokens.append(__UpperCAmelCase      )
              lowerCAmelCase__ :  Optional[Any]							      =					False
      out_string += self.sp_model.decode(__UpperCAmelCase      )
      return out_string.strip()




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		List[int]:
      lowerCAmelCase__ :  int							      =					[self.sep_token_id]
      lowerCAmelCase__ :  Dict							      =					[self.cls_token_id]
      if token_ids_a is None:
          return cls + token_ids_a + sep
      return cls + token_ids_a + sep + token_ids_a + sep




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None  ,__UpperCAmelCase = False      )    ->		List[int]:

      if already_has_special_tokens:
          return super().get_special_tokens_mask(
              token_ids_a=__UpperCAmelCase  ,token_ids_a=__UpperCAmelCase  ,already_has_special_tokens=__UpperCAmelCase      )

      if token_ids_a is not None:
          return [1] + ([0] * len(__UpperCAmelCase      )) + [1] + ([0] * len(__UpperCAmelCase      )) + [1]
      return [1] + ([0] * len(__UpperCAmelCase      )) + [1]




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		List[int]:
      lowerCAmelCase__ :  List[str]							      =					[self.sep_token_id]
      lowerCAmelCase__ :  Tuple							      =					[self.cls_token_id]

      if token_ids_a is None:
          return len(cls + token_ids_a + sep      ) * [0]
      return len(cls + token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1]






  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		Tuple[str]:
      if not os.path.isdir(__UpperCAmelCase      ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
          return
      lowerCAmelCase__ :  int							      =					os.path.join(
          __UpperCAmelCase  ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]      )

      if os.path.abspath(self.vocab_file      ) != os.path.abspath(__UpperCAmelCase      ) and os.path.isfile(self.vocab_file      ):
          copyfile(self.vocab_file  ,__UpperCAmelCase      )
      elif not os.path.isfile(self.vocab_file      ):
          with open(__UpperCAmelCase  ,"""wb"""      ) as fi:
              lowerCAmelCase__ :  List[Any]							      =					self.sp_model.serialized_model_proto()
              fi.write(__UpperCAmelCase      )

      return (out_vocab_file,)


