


'''simple docstring'''
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm import create_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import BitConfig, BitForImageClassification, BitImageProcessor
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_ :     Optional[int]								=      logging.get_logger(__name__)


def 			_lowercase       (				UpperCamelCase__       :					List[str]			):
 __A  :		str       =			'huggingface/label-files'
 __A  :		Optional[Any]       =			'imagenet-1k-id2label.json'
 __A  :		List[str]       =			json.load(open(hf_hub_download(UpperCamelCase__,				UpperCamelCase__,				repo_type='dataset'			),				'r'			)			)
 __A  :		Optional[Any]       =			{int(UpperCamelCase__			): v for k, v in idalabel.items()}
 __A  :		Optional[Any]       =			{v: k for k, v in idalabel.items()}

 __A  :		int       =			'std_conv' if 'bit' in model_name else False

 # note that when using BiT as backbone for ViT-hybrid checkpoints,
 # one needs to additionally set config.layer_type = "bottleneck", config.stem_type = "same",
 # config.conv_layer = "std_conv_same"
 __A  :		Optional[int]       =			BitConfig(
     conv_layer=UpperCamelCase__,				num_labels=1000,				idalabel=UpperCamelCase__,				labelaid=UpperCamelCase__,				)

 return config


def 			_lowercase       (				UpperCamelCase__       :					Union[str, Any]			):
 if "stem.conv" in name:
  __A  :		Optional[int]       =			name.replace('stem.conv',				'bit.embedder.convolution'			)
 if "blocks" in name:
  __A  :		int       =			name.replace('blocks',				'layers'			)
 if "head.fc" in name:
  __A  :		List[Any]       =			name.replace('head.fc',				'classifier.1'			)
 if name.startswith('norm'			):
  __A  :		Dict       =			'bit.' + name
 if "bit" not in name and "classifier" not in name:
  __A  :		Optional[int]       =			'bit.encoder.' + name

 return name


def 			_lowercase       (				):
 __A  :		str       =			'http://images.cocodataset.org/val2017/000000039769.jpg'
 __A  :		Union[str, Any]       =			Image.open(requests.get(UpperCamelCase__,				stream=UpperCamelCase__			).raw			)
 return im


@torch.no_grad()
def 			_lowercase       (				UpperCamelCase__       :					List[str],				UpperCamelCase__       :					str,				UpperCamelCase__       :					List[Any]=False			):
 __A  :		Optional[int]       =			get_config(UpperCamelCase__			)

 # load original model from timm
 __A  :		List[str]       =			create_model(UpperCamelCase__,				pretrained=UpperCamelCase__			)
 timm_model.eval()

 # load state_dict of original model
 __A  :		str       =			timm_model.state_dict()
 for key in state_dict.copy().keys():
  __A  :		int       =			state_dict.pop(UpperCamelCase__			)
  __A  :		Optional[Any]       =			val.squeeze() if 'head' in key else val

 # load HuggingFace model
 __A  :		Dict       =			BitForImageClassification(UpperCamelCase__			)
 model.eval()
 model.load_state_dict(UpperCamelCase__			)

 # create image processor
 __A  :		Tuple       =			create_transform(**resolve_data_config({},				model=UpperCamelCase__			)			)
 __A  :		Tuple       =			transform.transforms

 __A  :		List[Any]       =			{
     'bilinear': PILImageResampling.BILINEAR,
     'bicubic': PILImageResampling.BICUBIC,
     'nearest': PILImageResampling.NEAREST,
 }

 __A  :		List[Any]       =			BitImageProcessor(
     do_resize=UpperCamelCase__,				size={'shortest_edge': timm_transforms[0].size},				resample=pillow_resamplings[timm_transforms[0].interpolation.value],				do_center_crop=UpperCamelCase__,				crop_size={'height': timm_transforms[1].size[0], 'width': timm_transforms[1].size[1]},				do_normalize=UpperCamelCase__,				image_mean=timm_transforms[-1].mean.tolist(),				image_std=timm_transforms[-1].std.tolist(),				)

 __A  :		List[str]       =			prepare_img()
 __A  :		Union[str, Any]       =			transform(UpperCamelCase__			).unsqueeze(0			)
 __A  :		Optional[int]       =			processor(UpperCamelCase__,				return_tensors='pt'			).pixel_values

 # verify pixel values
 assert torch.allclose(UpperCamelCase__,				UpperCamelCase__			)

 # verify logits
 with torch.no_grad():
  __A  :		List[Any]       =			model(UpperCamelCase__			)
  __A  :		Any       =			outputs.logits

 print('Logits:',				logits[0, :3]			)
 print('Predicted class:',				model.config.idalabel[logits.argmax(-1			).item()]			)
 __A  :		Any       =			timm_model(UpperCamelCase__			)
 assert timm_logits.shape == outputs.logits.shape
 assert torch.allclose(UpperCamelCase__,				outputs.logits,				atol=1E-3			)
 print('Looks ok!'			)

 if pytorch_dump_folder_path is not None:
  Path(UpperCamelCase__			).mkdir(exist_ok=UpperCamelCase__			)
  print(f"""Saving model {model_name} and processor to {pytorch_dump_folder_path}"""			)
  model.save_pretrained(UpperCamelCase__			)
  processor.save_pretrained(UpperCamelCase__			)

 if push_to_hub:
  print(f"""Pushing model {model_name} and processor to the hub"""			)
  model.push_to_hub(f"""ybelkada/{model_name}"""			)
  processor.push_to_hub(f"""ybelkada/{model_name}"""			)


if __name__ == "__main__":
 UpperCAmelCase_ :     Any								=      argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--model_name',
     default='resnetv2_50x1_bitm',
     type=str,
     help='Name of the BiT timm 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 to the hub.',
 )

 UpperCAmelCase_ :     Optional[int]								=      parser.parse_args()
 convert_bit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''
from collections import defaultdict
from math import ceil, sqrt


def 			_lowercase       (				UpperCamelCase__       :					int = 1000000,				UpperCamelCase__       :					int = 10			):
 __A  :		defaultdict       =			defaultdict(UpperCamelCase__			)

 for outer_width in range(3,				(t_limit // 4) + 2			):
  if outer_width * outer_width > t_limit:
   __A  :		Optional[int]       =			max(
       ceil(sqrt(outer_width * outer_width - t_limit			)			),				1			)
  else:
   __A  :		Optional[Any]       =			1

  hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

  for hole_width in range(UpperCamelCase__,				outer_width - 1,				2			):
   count[outer_width * outer_width - hole_width * hole_width] += 1

 return sum(1 for n in count.values() if 1 <= n <= 10			)


if __name__ == "__main__":
 print(f'''{solution() = }''')






