




"""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   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		str:
 A__	=    [2, 2, 6, 2] if "tiny" in model_name else [2, 2, 18, 2]
 A__	=    True if "large" in model_name or "huge" in model_name else False
 A__	=    True if "large" in model_name or "huge" in model_name else False
 A__	=    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:
   A__	=    [3, 3, 3, 3]
   A__	=    [5, 5, 5, 5]
  elif "fl4" in model_name:
   A__	=    [4, 4, 4, 4]
   A__	=    [3, 3, 3, 3]

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

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

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

 A__	=    json.load(open(hf_hub_download(lowercase_      ,			lowercase_      ,			repo_type="dataset"						)      ,			"r"						)						)
 A__	=    {int(lowercase_						): v for k, v in idalabel.items()}
 A__	=    {v: k for k, v in idalabel.items()}

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

 if name == "norm.weight":
  A__	=    "layernorm.weight"
 if name == "norm.bias":
  A__	=    "layernorm.bias"

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

 return name







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

 A__	=    model_name_to_url[model_name]
 print("Checkpoint URL: "      ,			lowercase_						)
 A__	=    torch.hub.load_state_dict_from_url(lowercase_      ,			map_location="cpu"						)["model"]

 # rename keys
 for key in state_dict.copy().keys():
  A__	=    state_dict.pop(lowercase_						)
  A__	=    val

 A__	=    get_focalnet_config(lowercase_						)
 A__	=    FocalNetForImageClassification(lowercase_						)
 model.eval()

 # load state dict
 model.load_state_dict(lowercase_						)

 # verify conversion
 A__	=    "http://images.cocodataset.org/val2017/000000039769.jpg"

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

 A__	=    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]						),
     ]						)

 A__	=    image_transforms(lowercase_						).unsqueeze(0						)

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

 A__	=    model(**lowercase_						)

 A__	=    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":
  A__	=    torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]						)
 elif model_name == "focalnet-tiny-lrf":
  A__	=    torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]						)
 elif model_name == "focalnet-small":
  A__	=    torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]						)
 elif model_name == "focalnet-small-lrf":
  A__	=    torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]						)
 elif model_name == "focalnet-base":
  A__	=    torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]						)
 elif model_name == "focalnet-base-lrf":
  A__	=    torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]						)
 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__":
      SCREAMING_SNAKE_CASE    =	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.",
      )

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











"""simple docstring"""



import argparse
import os

import transformers

from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS
from .utils import logging


logging.set_verbosity_info()

SCREAMING_SNAKE_CASE    =	logging.get_logger(__name__)


SCREAMING_SNAKE_CASE    =	{name: getattr(transformers, name + "Fast") for name in SLOW_TO_FAST_CONVERTERS}







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			lowercase_      ,			lowercase_						)						->		Any:
 if tokenizer_name is not None and tokenizer_name not in TOKENIZER_CLASSES:
  raise ValueError(f"""Unrecognized tokenizer name, should be one of {list(TOKENIZER_CLASSES.keys()						)}."""						)

 if tokenizer_name is None:
  A__	=    TOKENIZER_CLASSES
 else:
  A__	=    {tokenizer_name: getattr(lowercase_      ,			tokenizer_name + "Fast"						)}

 logger.info(f"""Loading tokenizer classes: {tokenizer_names}"""						)

 for tokenizer_name in tokenizer_names:
  A__	=    TOKENIZER_CLASSES[tokenizer_name]

  A__	=    True
  if checkpoint_name is None:
   A__	=    list(tokenizer_class.max_model_input_sizes.keys()						)
  else:
   A__	=    [checkpoint_name]

  logger.info(f"""For tokenizer {tokenizer_class.__class__.__name__} loading checkpoints: {checkpoint_names}"""						)

  for checkpoint in checkpoint_names:
   logger.info(f"""Loading {tokenizer_class.__class__.__name__} {checkpoint}"""						)

   # Load tokenizer
   A__	=    tokenizer_class.from_pretrained(lowercase_      ,			force_download=lowercase_						)

   # Save fast tokenizer
   logger.info(f"""Save fast tokenizer to {dump_path} with prefix {checkpoint} add_prefix {add_prefix}"""						)

   # For organization names we create sub-directories
   if "/" in checkpoint:
    A__,				A__	=    checkpoint.split("/"						)
    A__	=    os.path.join(lowercase_      ,			lowercase_						)
   elif add_prefix:
    A__	=    checkpoint
    A__	=    dump_path
   else:
    A__	=    None
    A__	=    dump_path

   logger.info(f"""=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}"""						)

   if checkpoint in list(tokenizer.pretrained_vocab_files_map.values()						)[0]:
    A__	=    list(tokenizer.pretrained_vocab_files_map.values()						)[0][checkpoint]
    A__	=    file_path.split(lowercase_						)[-1][0]
    if next_char == "/":
     A__	=    os.path.join(lowercase_      ,			lowercase_						)
     A__	=    None

    logger.info(f"""=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}"""						)

   A__	=    tokenizer.save_pretrained(
       lowercase_      ,			legacy_format=lowercase_      ,			filename_prefix=lowercase_						)
   logger.info(f"""=> File names {file_names}"""						)

   for file_name in file_names:
    if not file_name.endswith("tokenizer.json"						):
     os.remove(lowercase_						)
     logger.info(f"""=> removing {file_name}"""						)


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE    =	argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--dump_path", default=None, type=str, required=True, help="Path to output generated fast tokenizer files."
      )
      parser.add_argument(
          "--tokenizer_name",
          default=None,
          type=str,
          help=(
              f'Optional tokenizer type selected in the list of {list(TOKENIZER_CLASSES.keys())}. If not given, will '
              "download and convert all the checkpoints from AWS."
          ),
      )
      parser.add_argument(
          "--checkpoint_name",
          default=None,
          type=str,
          help="Optional checkpoint name. If not given, will download and convert the canonical checkpoints from AWS.",
      )
      parser.add_argument(
          "--force_download",
          action="store_true",
          help="Re-download checkpoints.",
      )
      SCREAMING_SNAKE_CASE    =	parser.parse_args()

      convert_slow_checkpoint_to_fast(args.tokenizer_name, args.checkpoint_name, args.dump_path, args.force_download)





