




'''simple docstring'''

import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor
from transformers.utils import logging


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




def    snake_case__     (  )			->       List[Any]:





			'''simple docstring'''






			snake_case__          =						DetaConfig(
			    num_queries=900       ,      encoder_ffn_dim=2048       ,      decoder_ffn_dim=2048       ,      num_feature_levels=5       ,      assign_first_stage=a       ,      with_box_refine=a       ,      two_stage=a       ,      )

			# set labels
			snake_case__          =						91
			snake_case__          =						"""huggingface/label-files"""
			snake_case__          =						"""coco-detection-id2label.json"""
			snake_case__          =						json.load(open(cached_download(hf_hub_url(a       ,      a       ,      repo_type="""dataset"""	)	)       ,      """r"""	)	)
			snake_case__          =						{int(a	): v for k, v in idalabel.items()}
			snake_case__          =						idalabel
			snake_case__          =						{v: k for k, v in idalabel.items()}

			return config




def    snake_case__     (  a	)			->       int:





			'''simple docstring'''






			snake_case__          =						[]

			# stem
			# fmt: off
			rename_keys.append(("""backbone.0.body.conv1.weight""", """model.backbone.model.embedder.embedder.convolution.weight""")	)
			rename_keys.append(("""backbone.0.body.bn1.weight""", """model.backbone.model.embedder.embedder.normalization.weight""")	)
			rename_keys.append(("""backbone.0.body.bn1.bias""", """model.backbone.model.embedder.embedder.normalization.bias""")	)
			rename_keys.append(("""backbone.0.body.bn1.running_mean""", """model.backbone.model.embedder.embedder.normalization.running_mean""")	)
			rename_keys.append(("""backbone.0.body.bn1.running_var""", """model.backbone.model.embedder.embedder.normalization.running_var""")	)
			# stages
			for stage_idx in range(len(config.backbone_config.depths	)	):
						for layer_idx in range(config.backbone_config.depths[stage_idx]	):
									# shortcut
									if layer_idx == 0:
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var""",
												    )	)
									# 3 convs
									for i in range(3	):
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean""",
												    )	)
												rename_keys.append(
												    (
												        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var""",
												        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var""",
												    )	)
    # transformer encoder
			for i in range(config.encoder_layers	):
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight""", F"""model.encoder.layers.{i}.self_attn.sampling_offsets.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias""", F"""model.encoder.layers.{i}.self_attn.sampling_offsets.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.attention_weights.weight""", F"""model.encoder.layers.{i}.self_attn.attention_weights.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.attention_weights.bias""", F"""model.encoder.layers.{i}.self_attn.attention_weights.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.value_proj.weight""", F"""model.encoder.layers.{i}.self_attn.value_proj.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.value_proj.bias""", F"""model.encoder.layers.{i}.self_attn.value_proj.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.output_proj.weight""", F"""model.encoder.layers.{i}.self_attn.output_proj.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.output_proj.bias""", F"""model.encoder.layers.{i}.self_attn.output_proj.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.weight""", F"""model.encoder.layers.{i}.self_attn_layer_norm.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.bias""", F"""model.encoder.layers.{i}.self_attn_layer_norm.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.weight""", F"""model.encoder.layers.{i}.fc1.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.bias""", F"""model.encoder.layers.{i}.fc1.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.weight""", F"""model.encoder.layers.{i}.fc2.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.bias""", F"""model.encoder.layers.{i}.fc2.bias""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.weight""", F"""model.encoder.layers.{i}.final_layer_norm.weight""")	)
						rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.bias""", F"""model.encoder.layers.{i}.final_layer_norm.bias""")	)

			# transformer decoder
			for i in range(config.decoder_layers	):
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight""", F"""model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias""", F"""model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.attention_weights.weight""", F"""model.decoder.layers.{i}.encoder_attn.attention_weights.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.attention_weights.bias""", F"""model.decoder.layers.{i}.encoder_attn.attention_weights.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.value_proj.weight""", F"""model.decoder.layers.{i}.encoder_attn.value_proj.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.value_proj.bias""", F"""model.decoder.layers.{i}.encoder_attn.value_proj.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.output_proj.weight""", F"""model.decoder.layers.{i}.encoder_attn.output_proj.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.output_proj.bias""", F"""model.decoder.layers.{i}.encoder_attn.output_proj.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.weight""", F"""model.decoder.layers.{i}.encoder_attn_layer_norm.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.bias""", F"""model.decoder.layers.{i}.encoder_attn_layer_norm.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""", F"""model.decoder.layers.{i}.self_attn.out_proj.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", F"""model.decoder.layers.{i}.self_attn.out_proj.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.norm2.weight""", F"""model.decoder.layers.{i}.self_attn_layer_norm.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.norm2.bias""", F"""model.decoder.layers.{i}.self_attn_layer_norm.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.weight""", F"""model.decoder.layers.{i}.fc1.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.bias""", F"""model.decoder.layers.{i}.fc1.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.weight""", F"""model.decoder.layers.{i}.fc2.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.bias""", F"""model.decoder.layers.{i}.fc2.bias""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.weight""", F"""model.decoder.layers.{i}.final_layer_norm.weight""")	)
						rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.bias""", F"""model.decoder.layers.{i}.final_layer_norm.bias""")	)

			# fmt: on

			return rename_keys




def    snake_case__     (  a       ,      a       ,      a	)			->       List[Any]:





			'''simple docstring'''






			snake_case__          =						dct.pop(a	)
			snake_case__          =						val




def    snake_case__     (  a       ,      a	)			->       Optional[int]:





			'''simple docstring'''






			snake_case__          =						config.d_model
			for i in range(config.decoder_layers	):
						# read in weights + bias of input projection layer of self-attention
						snake_case__          =						state_dict.pop(F"""transformer.decoder.layers.{i}.self_attn.in_proj_weight"""	)
						snake_case__          =						state_dict.pop(F"""transformer.decoder.layers.{i}.self_attn.in_proj_bias"""	)
						# next, add query, keys and values (in that order) to the state dict
						snake_case__          =						in_proj_weight[:hidden_size, :]
						snake_case__          =						in_proj_bias[:hidden_size]
						snake_case__          =						in_proj_weight[
						    hidden_size : hidden_size * 2, :
						]
						snake_case__          =						in_proj_bias[hidden_size : hidden_size * 2]
						snake_case__          =						in_proj_weight[-hidden_size:, :]
						snake_case__          =						in_proj_bias[-hidden_size:]




def    snake_case__     (  )			->       Union[str, Any]:





			'''simple docstring'''






			snake_case__          =						"""http://images.cocodataset.org/val2017/000000039769.jpg"""
			snake_case__          =						Image.open(requests.get(a       ,      stream=a	).raw	)

			return im




@torch.no_grad()
def    snake_case__     (  a       ,      a       ,      a	)			->       str:





			'''simple docstring'''






			snake_case__          =						get_deta_config()

			# load original state dict
			if model_name == "deta-resnet-50":
						snake_case__          =						"""adet_checkpoint0011.pth"""
			elif model_name == "deta-resnet-50-24-epochs":
						snake_case__          =						"""adet_2x_checkpoint0023.pth"""
			else:
						raise ValueError(F"""Model name {model_name} not supported"""	)
			snake_case__          =						hf_hub_download(repo_id="""nielsr/deta-checkpoints"""       ,      filename=a	)
			snake_case__          =						torch.load(a       ,      map_location="""cpu"""	)["""model"""]

			# rename keys
			snake_case__          =						create_rename_keys(a	)
			for src, dest in rename_keys:
						rename_key(a       ,      a       ,      a	)
			read_in_decoder_q_k_v(a       ,      a	)

			# fix some prefixes
			for key in state_dict.copy().keys():
						if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
									snake_case__          =						state_dict.pop(a	)
									snake_case__          =						val
						if "input_proj" in key:
									snake_case__          =						state_dict.pop(a	)
									snake_case__          =						val
						if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
									snake_case__          =						state_dict.pop(a	)
									snake_case__          =						val

    # finally, create HuggingFace model and load state dict
			snake_case__          =						DetaForObjectDetection(a	)
			model.load_state_dict(a	)
			model.eval()

			snake_case__          =						"""cuda""" if torch.cuda.is_available() else """cpu"""
			model.to(a	)

			# load image processor
			snake_case__          =						DetaImageProcessor(format="""coco_detection"""	)

			# verify our conversion on image
			snake_case__          =						prepare_img()
			snake_case__          =						processor(images=a       ,      return_tensors="""pt"""	)
			snake_case__          =						encoding["""pixel_values"""]
			snake_case__          =						model(pixel_values.to(a	)	)

			# verify logits
			if model_name == "deta-resnet-50":
						snake_case__          =						torch.tensor(
						    [[-7.39_78, -2.54_06, -4.16_68], [-8.26_84, -3.99_33, -3.80_96], [-7.05_15, -3.79_73, -5.85_16]]	)
						snake_case__          =						torch.tensor([[0.50_43, 0.49_73, 0.99_98], [0.25_42, 0.54_89, 0.47_48], [0.54_90, 0.27_65, 0.05_70]]	)
			elif model_name == "deta-resnet-50-24-epochs":
						snake_case__          =						torch.tensor(
						    [[-7.16_88, -2.48_57, -4.86_69], [-7.86_30, -3.81_54, -4.26_74], [-7.27_30, -4.18_65, -5.53_23]]	)
						snake_case__          =						torch.tensor([[0.50_21, 0.49_71, 0.99_94], [0.25_46, 0.54_86, 0.47_31], [0.16_86, 0.19_86, 0.21_42]]	)

			assert torch.allclose(outputs.logits[0, :3, :3]       ,      expected_logits.to(a	)       ,      atol=1e-4	)
			assert torch.allclose(outputs.pred_boxes[0, :3, :3]       ,      expected_boxes.to(a	)       ,      atol=1e-4	)
			print("""Everything ok!"""	)

			if pytorch_dump_folder_path:
						# Save model and processor
						logger.info(F"""Saving PyTorch model and processor to {pytorch_dump_folder_path}..."""	)
						Path(a	).mkdir(exist_ok=a	)
						model.save_pretrained(a	)
						processor.save_pretrained(a	)

			# Push to hub
			if push_to_hub:
						print("""Pushing model and processor to hub..."""	)
						model.push_to_hub(F"""jozhang97/{model_name}"""	)
						processor.push_to_hub(F"""jozhang97/{model_name}"""	)


if __name__ == "__main__":
		a__         =     argparse.ArgumentParser()

		parser.add_argument(
		    '''--model_name''',
		    type=str,
		    default='''deta-resnet-50''',
		    choices=['''deta-resnet-50''', '''deta-resnet-50-24-epochs'''],
		    help='''Name of the model you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''',
		    default=None,
		    type=str,
		    help='''Path to the folder to output PyTorch model.''',
		)
		parser.add_argument(
		    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
		)
		a__         =     parser.parse_args()
		convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






"""simple docstring"""





import argparse
import logging

import datasets
from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
from tokenizers.models import Unigram
from tokenizers.trainers import UnigramTrainer

from transformers import AlbertTokenizerFast


lowerCAmelCase__        =      logging.getLogger(__name__)


def    a__       (       ):




		"""simple docstring"""


		UpperCamelCase =						argparse.ArgumentParser(description="Train a unigram tokenizer on the wikitext dataset."     )
		parser.add_argument(
		    "--dataset_name"	,						type=_SCREAMING_SNAKE_CASE	,						default="wikitext"	,						help="Name of the training. Explore datasets at: hf.co/datasets."	,						)
		parser.add_argument(
		    "--dataset_config"	,						type=_SCREAMING_SNAKE_CASE	,						default="wikitext-103-raw-v1"	,						help="Configuration name of the dataset."     )
		parser.add_argument(
		    "--batch_size"	,						type=_SCREAMING_SNAKE_CASE	,						default=1_000	,						help="Batch size during training."	,						)
		parser.add_argument(
		    "--vocab_size"	,						type=_SCREAMING_SNAKE_CASE	,						default=10_048	,						help="Size of the desired vocabulary."	,						)
		parser.add_argument(
		    "--limit"	,						default=_SCREAMING_SNAKE_CASE	,						type=_SCREAMING_SNAKE_CASE	,						help="Limit the number of shards (used for debugging)."	,						)
		parser.add_argument(
		    "--export_to_hub"	,						action="store_true"	,						)

		UpperCamelCase =						parser.parse_args()
		return args


def    a__       (       _SCREAMING_SNAKE_CASE     ):




		"""simple docstring"""


		UpperCamelCase =						datasets.load_dataset(args.dataset_name	,						args.dataset_config	,						split="train"     )

		if args.limit is not None:
				UpperCamelCase =						min(len(_SCREAMING_SNAKE_CASE     )	,						args.limit     )
				UpperCamelCase =						dataset.select(range(_SCREAMING_SNAKE_CASE     )     )
				logger.info(F"Limiting the dataset to {args.limit} entries."     )

		def batch_iterator():
				for i in range(0	,						len(_SCREAMING_SNAKE_CASE     )	,						args.batch_size     ):
						yield dataset[i : i + args.batch_size]["text"]

    # Prepare the tokenizer.
		UpperCamelCase =						Tokenizer(Unigram()     )
		UpperCamelCase =						normalizers.Sequence([normalizers.Replace("``"	,						"\""     ), normalizers.Replace("''"	,						"\""     )]     )
		UpperCamelCase =						pre_tokenizers.Metaspace()

		# Prepare the trainer.
		UpperCamelCase =						UnigramTrainer(
		    unk_token="<unk>"	,						special_tokens=["[CLS]", "[SEP]", "<unk>", "<pad>", "[MASK]"]	,						vocab_size=args.vocab_size	,						)

		logger.info("Training the tokenizer."     )
		tokenizer.train_from_iterator(batch_iterator()	,						trainer=_SCREAMING_SNAKE_CASE     )
		logger.info("Tokenizer training complete!"     )

		UpperCamelCase =						tokenizer.token_to_id("[CLS]"     )
		UpperCamelCase =						tokenizer.token_to_id("[SEP]"     )
		UpperCamelCase =						processors.TemplateProcessing(
		    single="[CLS]:0 $A:0 [SEP]:0"	,						pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1"	,						special_tokens=[
		        ("[CLS]", cls_token_id),
		        ("[SEP]", sep_token_id),
		    ]	,						)
		UpperCamelCase =						decoders.Metaspace()

		if args.export_to_hub:
				logger.info("Exporting the trained tokenzier to Hub."     )
				UpperCamelCase =						AlbertTokenizerFast(tokenizer_object=_SCREAMING_SNAKE_CASE     )
				new_tokenizer.push_to_hub("unigram-tokenizer-dataset"     )


if __name__ == "__main__":
						lowerCAmelCase__        =      parse_args()
						main(args)






