

'''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 transformers import YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


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






def 							_lowerCAmelCase (      __magic_name__	:    str	)							->     YolosConfig:
						lowercase   :       Optional[Any]								=YolosConfig()

						# size of the architecture
						if "yolos_ti" in yolos_name:
												lowercase   :       List[Any]								=192
												lowercase   :       Union[str, Any]								=768
												lowercase   :       List[Any]								=12
												lowercase   :       Tuple								=3
												lowercase   :       Tuple								=[800, 1333]
												lowercase   :       List[str]								=False
						elif yolos_name == "yolos_s_dWr":
												lowercase   :       Union[str, Any]								=330
												lowercase   :       Union[str, Any]								=14
												lowercase   :       int								=6
												lowercase   :       Tuple								=1320
						elif "yolos_s" in yolos_name:
												lowercase   :       Optional[int]								=384
												lowercase   :       int								=1536
												lowercase   :       Union[str, Any]								=12
												lowercase   :       List[str]								=6
						elif "yolos_b" in yolos_name:
												lowercase   :       Union[str, Any]								=[800, 1344]

						lowercase   :       Union[str, Any]								=91
						lowercase   :       int								='''huggingface/label-files'''
						lowercase   :       str								='''coco-detection-id2label.json'''
						lowercase   :       List[str]								=json.load(open(hf_hub_download(__magic_name__      ,   __magic_name__      ,   repo_type='''dataset'''	)      ,   '''r'''	)	)
						lowercase   :       List[Any]								={int(__magic_name__	): v for k, v in idalabel.items()}
						lowercase   :       Optional[int]								=idalabel
						lowercase   :       Any								={v: k for k, v in idalabel.items()}

						return config






def 							_lowerCAmelCase (      __magic_name__	:    dict      ,   __magic_name__	:    YolosConfig      ,   __magic_name__	:    bool = False	)							->     Optional[Any]:
						for i in range(config.num_hidden_layers	):
												# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
												lowercase   :       List[Any]								=state_dict.pop(f'''blocks.{i}.attn.qkv.weight'''	)
												lowercase   :       Tuple								=state_dict.pop(f'''blocks.{i}.attn.qkv.bias'''	)
												# next, add query, keys and values (in that order) to the state dict
												lowercase   :       Union[str, Any]								=in_proj_weight[: config.hidden_size, :]
												lowercase   :       Any								=in_proj_bias[: config.hidden_size]
												lowercase   :       int								=in_proj_weight[
												    config.hidden_size : config.hidden_size * 2, :
												]
												lowercase   :       int								=in_proj_bias[
												    config.hidden_size : config.hidden_size * 2
												]
												lowercase   :       Union[str, Any]								=in_proj_weight[-config.hidden_size :, :]
												lowercase   :       List[Any]								=in_proj_bias[-config.hidden_size :]






def 							_lowerCAmelCase (      __magic_name__	:    str	)							->     str:
						if "backbone" in name:
												lowercase   :       Optional[int]								=name.replace('''backbone'''      ,   '''vit'''	)
						if "cls_token" in name:
												lowercase   :       int								=name.replace('''cls_token'''      ,   '''embeddings.cls_token'''	)
						if "det_token" in name:
												lowercase   :       Tuple								=name.replace('''det_token'''      ,   '''embeddings.detection_tokens'''	)
						if "mid_pos_embed" in name:
												lowercase   :       Optional[int]								=name.replace('''mid_pos_embed'''      ,   '''encoder.mid_position_embeddings'''	)
						if "pos_embed" in name:
												lowercase   :       Union[str, Any]								=name.replace('''pos_embed'''      ,   '''embeddings.position_embeddings'''	)
						if "patch_embed.proj" in name:
												lowercase   :       Optional[Any]								=name.replace('''patch_embed.proj'''      ,   '''embeddings.patch_embeddings.projection'''	)
						if "blocks" in name:
												lowercase   :       Optional[Any]								=name.replace('''blocks'''      ,   '''encoder.layer'''	)
						if "attn.proj" in name:
												lowercase   :       Dict								=name.replace('''attn.proj'''      ,   '''attention.output.dense'''	)
						if "attn" in name:
												lowercase   :       str								=name.replace('''attn'''      ,   '''attention.self'''	)
						if "norm1" in name:
												lowercase   :       Dict								=name.replace('''norm1'''      ,   '''layernorm_before'''	)
						if "norm2" in name:
												lowercase   :       str								=name.replace('''norm2'''      ,   '''layernorm_after'''	)
						if "mlp.fc1" in name:
												lowercase   :       Optional[Any]								=name.replace('''mlp.fc1'''      ,   '''intermediate.dense'''	)
						if "mlp.fc2" in name:
												lowercase   :       int								=name.replace('''mlp.fc2'''      ,   '''output.dense'''	)
						if "class_embed" in name:
												lowercase   :       List[Any]								=name.replace('''class_embed'''      ,   '''class_labels_classifier'''	)
						if "bbox_embed" in name:
												lowercase   :       int								=name.replace('''bbox_embed'''      ,   '''bbox_predictor'''	)
						if "vit.norm" in name:
												lowercase   :       str								=name.replace('''vit.norm'''      ,   '''vit.layernorm'''	)

						return name






def 							_lowerCAmelCase (      __magic_name__	:    dict      ,   __magic_name__	:    YolosForObjectDetection	)							->     dict:
						for key in orig_state_dict.copy().keys():
												lowercase   :       List[str]								=orig_state_dict.pop(__magic_name__	)

												if "qkv" in key:
																		lowercase   :       Optional[int]								=key.split('''.'''	)
																		lowercase   :       int								=int(key_split[2]	)
																		lowercase   :       str								=model.vit.encoder.layer[layer_num].attention.attention.all_head_size
																		if "weight" in key:
																								lowercase   :       Any								=val[:dim, :]
																								lowercase   :       Tuple								=val[
																								    dim : dim * 2, :
																								]
																								lowercase   :       Optional[Any]								=val[-dim:, :]
																		else:
																								lowercase   :       List[Any]								=val[:dim]
																								lowercase   :       List[str]								=val[dim : dim * 2]
																								lowercase   :       Optional[Any]								=val[-dim:]
												else:
																		lowercase   :       Optional[Any]								=val

						return orig_state_dict






def 							_lowerCAmelCase (      )							->     torch.Tensor:
						lowercase   :       Optional[Any]								='''http://images.cocodataset.org/val2017/000000039769.jpg'''
						lowercase   :       Optional[int]								=Image.open(requests.get(__magic_name__      ,   stream=__magic_name__	).raw	)
						return im






@torch.no_grad()
def 							_lowerCAmelCase (      __magic_name__	:    str      ,   __magic_name__	:    str      ,   __magic_name__	:    str      ,   __magic_name__	:    bool = False	)							->     Union[str, Any]:
						lowercase   :       List[Any]								=get_yolos_config(__magic_name__	)

						# load original state_dict
						lowercase   :       Optional[int]								=torch.load(__magic_name__      ,   map_location='''cpu'''	)['''model''']

						# load 🤗 model
						lowercase   :       Any								=YolosForObjectDetection(__magic_name__	)
						model.eval()
						lowercase   :       Optional[Any]								=convert_state_dict(__magic_name__      ,   __magic_name__	)
						model.load_state_dict(__magic_name__	)

						# Check outputs on an image, prepared by YolosImageProcessor
						lowercase   :       Optional[int]								=800 if yolos_name != '''yolos_ti''' else 512
						lowercase   :       Any								=YolosImageProcessor(format='''coco_detection'''      ,   size=__magic_name__	)
						lowercase   :       str								=image_processor(images=prepare_img()      ,   return_tensors='''pt'''	)
						lowercase   :       List[Any]								=model(**__magic_name__	)
						lowercase		,  lowercase   :       List[str]								=outputs.logits, outputs.pred_boxes

						lowercase		,  lowercase   :       Any								=None, None
						if yolos_name == "yolos_ti":
												lowercase   :       Optional[Any]								=torch.tensor(
												    [[-3_9.5_0_2_2, -1_1.9_8_2_0, -1_7.6_8_8_8], [-2_9.9_5_7_4, -9.9_7_6_9, -1_7.7_6_9_1], [-4_2.3_2_8_1, -2_0.7_2_0_0, -3_0.6_2_9_4]]	)
												lowercase   :       int								=torch.tensor(
												    [[0.4_0_2_1, 0.0_8_3_6, 0.7_9_7_9], [0.0_1_8_4, 0.2_6_0_9, 0.0_3_6_4], [0.1_7_8_1, 0.2_0_0_4, 0.2_0_9_5]]	)
						elif yolos_name == "yolos_s_200_pre":
												lowercase   :       List[str]								=torch.tensor(
												    [[-2_4.0_2_4_8, -1_0.3_0_2_4, -1_4.8_2_9_0], [-4_2.0_3_9_2, -1_6.8_2_0_0, -2_7.4_3_3_4], [-2_7.2_7_4_3, -1_1.8_1_5_4, -1_8.7_1_4_8]]	)
												lowercase   :       List[Any]								=torch.tensor(
												    [[0.2_5_5_9, 0.5_4_5_5, 0.4_7_0_6], [0.2_9_8_9, 0.7_2_7_9, 0.1_8_7_5], [0.7_7_3_2, 0.4_0_1_7, 0.4_4_6_2]]	)
						elif yolos_name == "yolos_s_300_pre":
												lowercase   :       int								=torch.tensor(
												    [[-3_6.2_2_2_0, -1_4.4_3_8_5, -2_3.5_4_5_7], [-3_5.6_9_7_0, -1_4.7_5_8_3, -2_1.3_9_3_5], [-3_1.5_9_3_9, -1_3.6_0_4_2, -1_6.8_0_4_9]]	)
												lowercase   :       Dict								=torch.tensor(
												    [[0.7_6_1_4, 0.2_3_1_6, 0.4_7_2_8], [0.7_1_6_8, 0.4_4_9_5, 0.3_8_5_5], [0.4_9_9_6, 0.1_4_6_6, 0.9_9_9_6]]	)
						elif yolos_name == "yolos_s_dWr":
												lowercase   :       Optional[Any]								=torch.tensor(
												    [[-4_2.8_6_6_8, -2_4.1_0_4_9, -4_1.1_6_9_0], [-3_4.7_4_5_6, -1_4.1_2_7_4, -2_4.9_1_9_4], [-3_3.7_8_9_8, -1_2.1_9_4_6, -2_5.6_4_9_5]]	)
												lowercase   :       List[Any]								=torch.tensor(
												    [[0.5_5_8_7, 0.2_7_7_3, 0.0_6_0_5], [0.5_0_0_4, 0.3_0_1_4, 0.9_9_9_4], [0.4_9_9_9, 0.1_5_4_8, 0.9_9_9_4]]	)
						elif yolos_name == "yolos_base":
												lowercase   :       Union[str, Any]								=torch.tensor(
												    [[-4_0.6_0_6_4, -2_4.3_0_8_4, -3_2.6_4_4_7], [-5_5.1_9_9_0, -3_0.7_7_1_9, -3_5.5_8_7_7], [-5_1.4_3_1_1, -3_3.3_5_0_7, -3_5.6_4_6_2]]	)
												lowercase   :       List[str]								=torch.tensor(
												    [[0.5_5_5_5, 0.2_7_9_4, 0.0_6_5_5], [0.9_0_4_9, 0.2_6_6_4, 0.1_8_9_4], [0.9_1_8_3, 0.1_9_8_4, 0.1_6_3_5]]	)
						else:
												raise ValueError(f'''Unknown yolos_name: {yolos_name}'''	)

						assert torch.allclose(logits[0, :3, :3]      ,   __magic_name__      ,   atol=1E-4	)
						assert torch.allclose(pred_boxes[0, :3, :3]      ,   __magic_name__      ,   atol=1E-4	)

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

						if push_to_hub:
												lowercase   :       str								={
												    '''yolos_ti''': '''yolos-tiny''',
												    '''yolos_s_200_pre''': '''yolos-small''',
												    '''yolos_s_300_pre''': '''yolos-small-300''',
												    '''yolos_s_dWr''': '''yolos-small-dwr''',
												    '''yolos_base''': '''yolos-base''',
												}

												print('''Pushing to the hub...'''	)
												lowercase   :       Optional[Any]								=model_mapping[yolos_name]
												image_processor.push_to_hub(__magic_name__      ,   organization='''hustvl'''	)
												model.push_to_hub(__magic_name__      ,   organization='''hustvl'''	)


if __name__ == "__main__":
							UpperCamelCase_							=					argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--yolos_name""",
							    default="""yolos_s_200_pre""",
							    type=str,
							    help=(
							        """Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"""
							        """ 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."""
							    ),
							)
							parser.add_argument(
							    """--checkpoint_path""", default=None, type=str, help="""Path to the original state dict (.pth file)."""
							)
							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 or not to push the converted model to the 🤗 hub."""
							)

							UpperCamelCase_							=					parser.parse_args()
							convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)








'''simple docstring'''





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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


UpperCamelCase_							=					logging.get_logger(__name__)

UpperCamelCase_							=					"""▁"""

UpperCamelCase_							=					{"""vocab_file""": """sentencepiece.bpe.model"""}

UpperCamelCase_							=					{
    """vocab_file""": {
        """facebook/xglm-564M""": """https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model""",
    }
}

UpperCamelCase_							=					{
    """facebook/xglm-564M""": 2048,
}





class 		__SCREAMING_SNAKE_CASE       ( lowercase__			):
				lowerCamelCase_							=						VOCAB_FILES_NAMES
				lowerCamelCase_							=						PRETRAINED_VOCAB_FILES_MAP
				lowerCamelCase_							=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				lowerCamelCase_							=						['input_ids', 'attention_mask']







				def __init__(					self    :		Dict		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		List[Any]="<s>"		,				UpperCAmelCase__    :		int="</s>"		,				UpperCAmelCase__    :		Optional[Any]="</s>"		,				UpperCAmelCase__    :		Optional[Any]="<s>"		,				UpperCAmelCase__    :		Any="<unk>"		,				UpperCAmelCase__    :		Any="<pad>"		,				UpperCAmelCase__    :		Optional[Dict[str, Any]] = None		,				**UpperCAmelCase__    :		Optional[int]		,				):
										'''simple docstring'''
										lowercase   :       int								={} if sp_model_kwargs is None else sp_model_kwargs

										# Compatibility with the original tokenizer
										lowercase   :       Optional[Any]								=7
										lowercase   :       Optional[int]								=[F'''<madeupword{i}>''' for i in range(self.num_madeup_words      )]

										lowercase   :       List[Any]								=kwargs.get('''additional_special_tokens'''		,				[]      )
										kwargs["additional_special_tokens"] += [
										    word for word in madeup_words if word not in kwargs["additional_special_tokens"]
										]

										super().__init__(
										    bos_token=UpperCAmelCase__		,				eos_token=UpperCAmelCase__		,				unk_token=UpperCAmelCase__		,				sep_token=UpperCAmelCase__		,				cls_token=UpperCAmelCase__		,				pad_token=UpperCAmelCase__		,				sp_model_kwargs=self.sp_model_kwargs		,				**UpperCAmelCase__		,				)

										lowercase   :       List[str]								=spm.SentencePieceProcessor(**self.sp_model_kwargs      )
										self.sp_model.Load(str(UpperCAmelCase__      )      )
										lowercase   :       List[Any]								=vocab_file

										# Original fairseq vocab and spm vocab must be "aligned":
										# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
										# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
										# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
										# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

										# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
										lowercase   :       Union[str, Any]								=1

										# Mimic fairseq token-to-id alignment for the first 4 token
										lowercase   :       List[str]								={'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

										lowercase   :       str								=len(self.sp_model      )
										lowercase   :       List[Any]								={F'''<madeupword{i}>''': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words      )}
										self.fairseq_tokens_to_ids.update(UpperCAmelCase__      )

										lowercase   :       int								={v: k for k, v in self.fairseq_tokens_to_ids.items()}







				def __getstate__(					self    :		int      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=self.__dict__.copy()
										lowercase   :       List[Any]								=None
										lowercase   :       Tuple								=self.sp_model.serialized_model_proto()
										return state







				def __setstate__(					self    :		Optional[Any]		,				UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										lowercase   :       int								=d

										# for backward compatibility
										if not hasattr(self		,				'''sp_model_kwargs'''      ):
																lowercase   :       Optional[int]								={}

										lowercase   :       List[Any]								=spm.SentencePieceProcessor(**self.sp_model_kwargs      )
										self.sp_model.LoadFromSerializedProto(self.sp_model_proto      )







				def    lowerCamelCase_  (					self    :		str		,				UpperCAmelCase__    :		List[int]		,				UpperCAmelCase__    :		Optional[List[int]] = None      ):
										'''simple docstring'''

										if token_ids_a is None:
																return [self.sep_token_id] + token_ids_a
										lowercase   :       List[Any]								=[self.sep_token_id]
										return sep + token_ids_a + sep + sep + token_ids_a







				def    lowerCamelCase_  (					self    :		Any		,				UpperCAmelCase__    :		List[int]		,				UpperCAmelCase__    :		Optional[List[int]] = None		,				UpperCAmelCase__    :		bool = False      ):
										'''simple docstring'''

										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 None:
																return [1] + ([0] * len(UpperCAmelCase__      ))
										return [1] + ([0] * len(UpperCAmelCase__      )) + [1, 1] + ([0] * len(UpperCAmelCase__      ))







				def    lowerCamelCase_  (					self    :		Dict		,				UpperCAmelCase__    :		List[int]		,				UpperCAmelCase__    :		Optional[List[int]] = None      ):
										'''simple docstring'''
										lowercase   :       int								=[self.sep_token_id]

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







				@property
				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										return len(self.sp_model      ) + self.fairseq_offset + self.num_madeup_words







				def    lowerCamelCase_  (					self    :		List[str]      ):
										'''simple docstring'''
										lowercase   :       int								={self.convert_ids_to_tokens(UpperCAmelCase__      ): i for i in range(self.vocab_size      )}
										vocab.update(self.added_tokens_encoder      )
										return vocab







				def    lowerCamelCase_  (					self    :		Tuple		,				UpperCAmelCase__    :		str      ):
										'''simple docstring'''
										return self.sp_model.encode(UpperCAmelCase__		,				out_type=UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Dict		,				UpperCAmelCase__    :		Optional[Any]      ):
										'''simple docstring'''
										if token in self.fairseq_tokens_to_ids:
																return self.fairseq_tokens_to_ids[token]
										lowercase   :       List[str]								=self.sp_model.PieceToId(UpperCAmelCase__      )

										# Need to return unknown token if the SP model returned 0
										return spm_id + self.fairseq_offset if spm_id else self.unk_token_id







				def    lowerCamelCase_  (					self    :		Optional[Any]		,				UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										if index in self.fairseq_ids_to_tokens:
																return self.fairseq_ids_to_tokens[index]
										return self.sp_model.IdToPiece(index - self.fairseq_offset      )







				def    lowerCamelCase_  (					self    :		List[str]		,				UpperCAmelCase__    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       Dict								=''''''.join(UpperCAmelCase__      ).replace(UpperCAmelCase__		,				''' '''      ).strip()
										return out_string







				def    lowerCamelCase_  (					self    :		Optional[int]		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		Optional[str] = None      ):
										'''simple docstring'''
										if not os.path.isdir(UpperCAmelCase__      ):
																logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''      )
																return
										lowercase   :       Dict								=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:
																						lowercase   :       Optional[int]								=self.sp_model.serialized_model_proto()
																						fi.write(UpperCAmelCase__      )

										return (out_vocab_file,)





