





import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
__a  :					Any          =  logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
__a  :					Tuple          =  []
for i in range(6):
		# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
		rename_keys.append(
		    (F'''transformer.encoder.layers.{i}.self_attn.out_proj.weight''', F'''encoder.layers.{i}.self_attn.out_proj.weight''')
		)
		rename_keys.append(
		    (F'''transformer.encoder.layers.{i}.self_attn.out_proj.bias''', F'''encoder.layers.{i}.self_attn.out_proj.bias''')
		)
		rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.weight''', F'''encoder.layers.{i}.fc1.weight'''))
		rename_keys.append((F'''transformer.encoder.layers.{i}.linear1.bias''', F'''encoder.layers.{i}.fc1.bias'''))
		rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.weight''', F'''encoder.layers.{i}.fc2.weight'''))
		rename_keys.append((F'''transformer.encoder.layers.{i}.linear2.bias''', F'''encoder.layers.{i}.fc2.bias'''))
		rename_keys.append(
		    (F'''transformer.encoder.layers.{i}.norm1.weight''', F'''encoder.layers.{i}.self_attn_layer_norm.weight''')
		)
		rename_keys.append((F'''transformer.encoder.layers.{i}.norm1.bias''', F'''encoder.layers.{i}.self_attn_layer_norm.bias'''))
		rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.weight''', F'''encoder.layers.{i}.final_layer_norm.weight'''))
		rename_keys.append((F'''transformer.encoder.layers.{i}.norm2.bias''', F'''encoder.layers.{i}.final_layer_norm.bias'''))
		# decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
		rename_keys.append(
		    (F'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', F'''decoder.layers.{i}.self_attn.out_proj.weight''')
		)
		rename_keys.append(
		    (F'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', F'''decoder.layers.{i}.self_attn.out_proj.bias''')
		)
		rename_keys.append(
		    (
		        F'''transformer.decoder.layers.{i}.multihead_attn.out_proj.weight''',
		        F'''decoder.layers.{i}.encoder_attn.out_proj.weight''',
		    )
		)
		rename_keys.append(
		    (
		        F'''transformer.decoder.layers.{i}.multihead_attn.out_proj.bias''',
		        F'''decoder.layers.{i}.encoder_attn.out_proj.bias''',
		    )
		)
		rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.weight''', F'''decoder.layers.{i}.fc1.weight'''))
		rename_keys.append((F'''transformer.decoder.layers.{i}.linear1.bias''', F'''decoder.layers.{i}.fc1.bias'''))
		rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.weight''', F'''decoder.layers.{i}.fc2.weight'''))
		rename_keys.append((F'''transformer.decoder.layers.{i}.linear2.bias''', F'''decoder.layers.{i}.fc2.bias'''))
		rename_keys.append(
		    (F'''transformer.decoder.layers.{i}.norm1.weight''', F'''decoder.layers.{i}.self_attn_layer_norm.weight''')
		)
		rename_keys.append((F'''transformer.decoder.layers.{i}.norm1.bias''', F'''decoder.layers.{i}.self_attn_layer_norm.bias'''))
		rename_keys.append(
		    (F'''transformer.decoder.layers.{i}.norm2.weight''', F'''decoder.layers.{i}.encoder_attn_layer_norm.weight''')
		)
		rename_keys.append(
		    (F'''transformer.decoder.layers.{i}.norm2.bias''', F'''decoder.layers.{i}.encoder_attn_layer_norm.bias''')
		)
		rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.weight''', F'''decoder.layers.{i}.final_layer_norm.weight'''))
		rename_keys.append((F'''transformer.decoder.layers.{i}.norm3.bias''', F'''decoder.layers.{i}.final_layer_norm.bias'''))

# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ("""input_proj.weight""", """input_projection.weight"""),
        ("""input_proj.bias""", """input_projection.bias"""),
        ("""query_embed.weight""", """query_position_embeddings.weight"""),
        ("""transformer.decoder.norm.weight""", """decoder.layernorm.weight"""),
        ("""transformer.decoder.norm.bias""", """decoder.layernorm.bias"""),
        ("""class_embed.weight""", """class_labels_classifier.weight"""),
        ("""class_embed.bias""", """class_labels_classifier.bias"""),
        ("""bbox_embed.layers.0.weight""", """bbox_predictor.layers.0.weight"""),
        ("""bbox_embed.layers.0.bias""", """bbox_predictor.layers.0.bias"""),
        ("""bbox_embed.layers.1.weight""", """bbox_predictor.layers.1.weight"""),
        ("""bbox_embed.layers.1.bias""", """bbox_predictor.layers.1.bias"""),
        ("""bbox_embed.layers.2.weight""", """bbox_predictor.layers.2.weight"""),
        ("""bbox_embed.layers.2.bias""", """bbox_predictor.layers.2.bias"""),
    ]
)





def     UpperCAmelCase						(   lowercase							,      lowercase							,      lowercase							):







				"""simple docstring"""



				__lowercase			      =							state_dict.pop(__lowercase							)
				__lowercase			      =							val





def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							OrderedDict()
				for key, value in state_dict.items():
								if "backbone.0.body" in key:
												__lowercase			      =							key.replace('''backbone.0.body'''							,      '''backbone.conv_encoder.model'''							)
												__lowercase			      =							value
								else:
												__lowercase			      =							value

				return new_state_dict





def     UpperCAmelCase						(   lowercase							,      lowercase=False							):







				"""simple docstring"""



				__lowercase			      =							''''''
				if is_panoptic:
								__lowercase			      =							'''detr.'''

				# first: transformer encoder
				for i in range(6							):
								# read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
								__lowercase			      =							state_dict.pop(F"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"							)
								__lowercase			      =							state_dict.pop(F"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"							)
								# next, add query, keys and values (in that order) to the state dict
								__lowercase			      =							in_proj_weight[:256, :]
								__lowercase			      =							in_proj_bias[:256]
								__lowercase			      =							in_proj_weight[256:512, :]
								__lowercase			      =							in_proj_bias[256:512]
								__lowercase			      =							in_proj_weight[-256:, :]
								__lowercase			      =							in_proj_bias[-256:]
				# next: transformer decoder (which is a bit more complex because it also includes cross-attention)
				for i in range(6							):
								# read in weights + bias of input projection layer of self-attention
								__lowercase			      =							state_dict.pop(F"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"							)
								__lowercase			      =							state_dict.pop(F"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"							)
								# next, add query, keys and values (in that order) to the state dict
								__lowercase			      =							in_proj_weight[:256, :]
								__lowercase			      =							in_proj_bias[:256]
								__lowercase			      =							in_proj_weight[256:512, :]
								__lowercase			      =							in_proj_bias[256:512]
								__lowercase			      =							in_proj_weight[-256:, :]
								__lowercase			      =							in_proj_bias[-256:]
								# read in weights + bias of input projection layer of cross-attention
								__lowercase			      =							state_dict.pop(
								    F"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"							)
								__lowercase			      =							state_dict.pop(F"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"							)
								# next, add query, keys and values (in that order) of cross-attention to the state dict
								__lowercase			      =							in_proj_weight_cross_attn[:256, :]
								__lowercase			      =							in_proj_bias_cross_attn[:256]
								__lowercase			      =							in_proj_weight_cross_attn[256:512, :]
								__lowercase			      =							in_proj_bias_cross_attn[256:512]
								__lowercase			      =							in_proj_weight_cross_attn[-256:, :]
								__lowercase			      =							in_proj_bias_cross_attn[-256:]





def     UpperCAmelCase						(   ):







				"""simple docstring"""



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

				return im





@torch.no_grad()
def     UpperCAmelCase						(   lowercase							,      lowercase							):







				"""simple docstring"""



				__lowercase			      =							DetrConfig()
				# set backbone and dilation attributes
				if "resnet101" in model_name:
								__lowercase			      =							'''resnet101'''
				if "dc5" in model_name:
								__lowercase			      =							True
				__lowercase			      =							'''panoptic''' in model_name
				if is_panoptic:
								__lowercase			      =							250
				else:
								__lowercase			      =							91
								__lowercase			      =							'''huggingface/label-files'''
								__lowercase			      =							'''coco-detection-id2label.json'''
								__lowercase			      =							json.load(open(hf_hub_download(__lowercase							,      __lowercase							,      repo_type='''dataset'''							)							,      '''r'''							)							)
								__lowercase			      =							{int(__lowercase							): v for k, v in idalabel.items()}
								__lowercase			      =							idalabel
								__lowercase			      =							{v: k for k, v in idalabel.items()}

				# load image processor
				__lowercase			      =							'''coco_panoptic''' if is_panoptic else '''coco_detection'''
				__lowercase			      =							DetrImageProcessor(format=__lowercase							)

				# prepare image
				__lowercase			      =							prepare_img()
				__lowercase			      =							image_processor(images=__lowercase							,      return_tensors='''pt'''							)
				__lowercase			      =							encoding['''pixel_values''']

				logger.info(F"Converting model {model_name}..."							)

				# load original model from torch hub
				__lowercase			      =							torch.hub.load('''facebookresearch/detr'''							,      __lowercase							,      pretrained=__lowercase							).eval()
				__lowercase			      =							detr.state_dict()
				# rename keys
				for src, dest in rename_keys:
								if is_panoptic:
												__lowercase			      =							'''detr.''' + src
								rename_key(__lowercase							,      __lowercase							,      __lowercase							)
				__lowercase			      =							rename_backbone_keys(__lowercase							)
				# query, key and value matrices need special treatment
				read_in_q_k_v(__lowercase							,      is_panoptic=__lowercase							)
				# important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
				__lowercase			      =							'''detr.model.''' if is_panoptic else '''model.'''
				for key in state_dict.copy().keys():
								if is_panoptic:
												if (
												    key.startswith('''detr'''							)
												    and not key.startswith('''class_labels_classifier'''							)
												    and not key.startswith('''bbox_predictor'''							)
												):
																__lowercase			      =							state_dict.pop(__lowercase							)
																__lowercase			      =							val
												elif "class_labels_classifier" in key or "bbox_predictor" in key:
																__lowercase			      =							state_dict.pop(__lowercase							)
																__lowercase			      =							val
												elif key.startswith('''bbox_attention'''							) or key.startswith('''mask_head'''							):
																continue
												else:
																__lowercase			      =							state_dict.pop(__lowercase							)
																__lowercase			      =							val
								else:
												if not key.startswith('''class_labels_classifier'''							) and not key.startswith('''bbox_predictor'''							):
																__lowercase			      =							state_dict.pop(__lowercase							)
																__lowercase			      =							val
    # finally, create HuggingFace model and load state dict
				__lowercase			      =							DetrForSegmentation(__lowercase							) if is_panoptic else DetrForObjectDetection(__lowercase							)
				model.load_state_dict(__lowercase							)
				model.eval()
				# verify our conversion
				__lowercase			      =							detr(__lowercase							)
				__lowercase			      =							model(__lowercase							)
				assert torch.allclose(outputs.logits							,      original_outputs['''pred_logits''']							,      atol=1E-4							)
				assert torch.allclose(outputs.pred_boxes							,      original_outputs['''pred_boxes''']							,      atol=1E-4							)
				if is_panoptic:
								assert torch.allclose(outputs.pred_masks							,      original_outputs['''pred_masks''']							,      atol=1E-4							)

				# Save model and image processor
				logger.info(F"Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."							)
				Path(__lowercase							).mkdir(exist_ok=__lowercase							)
				model.save_pretrained(__lowercase							)
				image_processor.save_pretrained(__lowercase							)


if __name__ == "__main__":
		__a  :					int          =  argparse.ArgumentParser()

		parser.add_argument(
		    """--model_name""", default="""detr_resnet50""", type=str, help="""Name of the DETR 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."""
		)
		__a  :					Optional[Any]          =  parser.parse_args()
		convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)
import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__a  :					int          =  5_0_0_0_0_0

__a   ,						__a  :					Any          =  os.path.split(__file__)
__a  :					Dict          =  os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.select(range(0							,      len(lowercase							)							,      2							)							)





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.sort('''numbers'''							)





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.shuffle()





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.train_test_split(0.1							)





@get_duration
def     UpperCAmelCase						(   lowercase							,      lowercase=10							):







				"""simple docstring"""



				for shard_id in range(lowercase							):
								__lowercase			      =							dataset.shard(lowercase							,      lowercase							)





def     UpperCAmelCase						(   ):







				"""simple docstring"""



				__lowercase			      =							{'''num examples''': SPEED_TEST_N_EXAMPLES}
				__lowercase			      =							(select, sort, shuffle, train_test_split, shard)
				with tempfile.TemporaryDirectory() as tmp_dir:
								print('''generating dataset'''							)
								__lowercase			      =							datasets.Features({'''text''': datasets.Value('''string'''							), '''numbers''': datasets.Value('''float32'''							)}							)
								__lowercase			      =							generate_example_dataset(
								    os.path.join(lowercase							,      '''dataset.arrow'''							)							,      lowercase							,      num_examples=lowercase							)
								print('''Functions'''							)
								for func in functions:
												print(func.__name__							)
												__lowercase			      =							func(lowercase							)

				with open(lowercase							,      '''wb'''							) as f:
								f.write(json.dumps(lowercase							).encode('''utf-8'''							)							)


if __name__ == "__main__":  # useful to run the profiler
		benchmark_indices_mapping()