





"""simple docstring"""




from copy import deepcopy
from unittest.case import TestCase

import pytest

from datasets.arrow_dataset import Dataset
from datasets.features import Audio, ClassLabel, Features, Image, Sequence, Value
from datasets.info import DatasetInfo
from datasets.tasks import (
    AudioClassification,
    AutomaticSpeechRecognition,
    ImageClassification,
    LanguageModeling,
    QuestionAnsweringExtractive,
    Summarization,
    TextClassification,
    task_template_from_dict,
)
from datasets.utils.py_utils import asdict


lowerCamelCase_	:		List[Any]         =   {
    'id': '5733be284776f41900661182',
    'title': 'University_of_Notre_Dame',
    'context': 'Architecturally, the school has a Catholic character. Atop the Main Building\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.',
    'question': 'To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?',
    'answers': {'text': ['Saint Bernadette Soubirous'], 'answer_start': [5_15]},
}





@pytest.mark.parametrize(
    'task_cls'					,   [
        AudioClassification,
        AutomaticSpeechRecognition,
        ImageClassification,
        LanguageModeling,
        QuestionAnsweringExtractive,
        Summarization,
        TextClassification,
    ]					,   )
def 		UpperCAmelCase__						(	_UpperCAmelCase   ):






  """simple docstring"""





  A_    :			str									=    task_cls()
  A_    :			Optional[int]									=    asdict(_UpperCAmelCase   )
  A_    :			Union[str, Any]									=    task_template_from_dict(_UpperCAmelCase   )
  assert task == reloaded

class 					_UpperCAmelCase    :



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[str]									=    LanguageModeling(text_column='input_text'    )
        assert {"input_text": "text"} == task.column_mapping

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			int									=    Features({'text': Value('string'    )}    )
        A_    :			Optional[int]									=    {'text_column': 'input_text'}
        A_    :			List[str]									=    LanguageModeling.from_dict(snake_case_    )
        assert "language-modeling" == task.task
        assert input_schema == task.input_schema


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    sorted(['pos', 'neg']    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			int									=    TextClassification(text_column='input_text'   , label_column='input_label'    )
        self.assertDictEqual({'input_text': 'text', 'input_label': 'labels'}   , task.column_mapping    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    Features({'text': Value('string'    )}    )
        # Labels are cast to tuple during `TextClassification.__post_init__`, so we do the same here
        A_    :			List[str]									=    Features({'labels': ClassLabel}    )
        A_    :			List[str]									=    {'text_column': 'input_text', 'label_column': 'input_labels'}
        A_    :			List[str]									=    TextClassification.from_dict(snake_case_    )
        self.assertEqual('text-classification'   , task.task    )
        self.assertEqual(snake_case_   , task.input_schema    )
        self.assertEqual(snake_case_   , task.label_schema    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    TextClassification(text_column='input_text'   , label_column='input_label'    )
        self.assertEqual(task.label_schema['labels']   , snake_case_    )
        A_    :			Optional[int]									=    task.align_with_features(Features({'input_label': ClassLabel(names=self.labels    )}    )    )
        self.assertEqual(task.label_schema['labels']   , ClassLabel(names=self.labels    )    )


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Dict									=    QuestionAnsweringExtractive(
            context_column='input_context'   , question_column='input_question'   , answers_column='input_answers'    )
        self.assertDictEqual(
            {'input_context': 'context', 'input_question': 'question', 'input_answers': 'answers'}   , task.column_mapping    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    Features({'question': Value('string'    ), 'context': Value('string'    )}    )
        A_    :			Tuple									=    Features(
            {
                'answers': Sequence(
                    {
                        'text': Value('string'    ),
                        'answer_start': Value('int32'    ),
                    }    )
            }    )
        A_    :			int									=    {
            'context_column': 'input_input_context',
            'question_column': 'input_question',
            'answers_column': 'input_answers',
        }
        A_    :			Any									=    QuestionAnsweringExtractive.from_dict(snake_case_    )
        self.assertEqual('question-answering-extractive'   , task.task    )
        self.assertEqual(snake_case_   , task.input_schema    )
        self.assertEqual(snake_case_   , task.label_schema    )


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    Summarization(text_column='input_text'   , summary_column='input_summary'    )
        self.assertDictEqual({'input_text': 'text', 'input_summary': 'summary'}   , task.column_mapping    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[Any]									=    Features({'text': Value('string'    )}    )
        A_    :			Any									=    Features({'summary': Value('string'    )}    )
        A_    :			List[str]									=    {'text_column': 'input_text', 'summary_column': 'input_summary'}
        A_    :			Tuple									=    Summarization.from_dict(snake_case_    )
        self.assertEqual('summarization'   , task.task    )
        self.assertEqual(snake_case_   , task.input_schema    )
        self.assertEqual(snake_case_   , task.label_schema    )


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    AutomaticSpeechRecognition(audio_column='input_audio'   , transcription_column='input_transcription'    )
        self.assertDictEqual({'input_audio': 'audio', 'input_transcription': 'transcription'}   , task.column_mapping    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Tuple									=    Features({'audio': Audio()}    )
        A_    :			List[Any]									=    Features({'transcription': Value('string'    )}    )
        A_    :			Tuple									=    {
            'audio_column': 'input_audio',
            'transcription_column': 'input_transcription',
        }
        A_    :			Any									=    AutomaticSpeechRecognition.from_dict(snake_case_    )
        self.assertEqual('automatic-speech-recognition'   , task.task    )
        self.assertEqual(snake_case_   , task.input_schema    )
        self.assertEqual(snake_case_   , task.label_schema    )


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Dict									=    sorted(['pos', 'neg']    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			str									=    AudioClassification(audio_column='input_audio'   , label_column='input_label'    )
        self.assertDictEqual({'input_audio': 'audio', 'input_label': 'labels'}   , task.column_mapping    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Tuple									=    Features({'audio': Audio()}    )
        A_    :			Optional[Any]									=    Features({'labels': ClassLabel}    )
        A_    :			Optional[Any]									=    {
            'audio_column': 'input_image',
            'label_column': 'input_label',
        }
        A_    :			Optional[Any]									=    AudioClassification.from_dict(snake_case_    )
        self.assertEqual('audio-classification'   , task.task    )
        self.assertEqual(snake_case_   , task.input_schema    )
        self.assertEqual(snake_case_   , task.label_schema    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Dict									=    AudioClassification(audio_column='input_audio'   , label_column='input_label'    )
        self.assertEqual(task.label_schema['labels']   , snake_case_    )
        A_    :			Optional[Any]									=    task.align_with_features(Features({'input_label': ClassLabel(names=self.labels    )}    )    )
        self.assertEqual(task.label_schema['labels']   , ClassLabel(names=self.labels    )    )


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			str									=    sorted(['pos', 'neg']    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Any									=    ImageClassification(image_column='input_image'   , label_column='input_label'    )
        self.assertDictEqual({'input_image': 'image', 'input_label': 'labels'}   , task.column_mapping    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			int									=    Features({'image': Image()}    )
        A_    :			List[Any]									=    Features({'labels': ClassLabel}    )
        A_    :			Any									=    {
            'image_column': 'input_image',
            'label_column': 'input_label',
        }
        A_    :			int									=    ImageClassification.from_dict(snake_case_    )
        self.assertEqual('image-classification'   , task.task    )
        self.assertEqual(snake_case_   , task.input_schema    )
        self.assertEqual(snake_case_   , task.label_schema    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Tuple									=    ImageClassification(image_column='input_image'   , label_column='input_label'    )
        self.assertEqual(task.label_schema['labels']   , snake_case_    )
        A_    :			Dict									=    task.align_with_features(Features({'input_label': ClassLabel(names=self.labels    )}    )    )
        self.assertEqual(task.label_schema['labels']   , ClassLabel(names=self.labels    )    )


class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[Any]									=    DatasetInfo(task_templates=QuestionAnsweringExtractive()    )
        A_    :			Dict									=    Dataset.from_dict({k: [v] for k, v in SAMPLE_QUESTION_ANSWERING_EXTRACTIVE.items()}   , info=snake_case_    )
        assert isinstance(dataset.info.task_templates   , snake_case_    )
        assert len(dataset.info.task_templates    ) == 1

        def keep_task(snake_case_    ):
          return x

        def dont_keep_task(snake_case_    ):
          A_    :			int									=    deepcopy(snake_case_    )
          A_    :			Optional[Any]									=    0
          return out

        A_    :			List[str]									=    dataset.map(snake_case_    )
        assert mapped_dataset.info.task_templates == dataset.info.task_templates
        # reload from cache
        A_    :			Tuple									=    dataset.map(snake_case_    )
        assert mapped_dataset.info.task_templates == dataset.info.task_templates

        A_    :			str									=    dataset.map(snake_case_    )
        assert mapped_dataset.info.task_templates == []
        # reload from cache
        A_    :			Optional[Any]									=    dataset.map(snake_case_    )
        assert mapped_dataset.info.task_templates == []

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[int]									=    Features({'text': Value('string'    ), 'label': ClassLabel(names=('pos', 'neg')    )}    )
        A_    :			str									=    TextClassification(text_column='text'   , label_column='label'    )
        A_    :			List[str]									=    DatasetInfo(features=snake_case_   , task_templates=snake_case_    )
        A_    :			Tuple									=    Dataset.from_dict({'text': ['A sentence.'], 'label': ['pos']}   , info=snake_case_    )

        def process(snake_case_    ):
          return example

        A_    :			Optional[Any]									=    dataset.remove_columns('label'    )
        A_    :			str									=    modified_dataset.map(snake_case_    )
        assert mapped_dataset.info.task_templates == []


"""simple docstring"""

import argparse
import json
from pathlib import Path

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

from transformers import DeformableDetrConfig, DeformableDetrForObjectDetection, DeformableDetrImageProcessor
from transformers.utils import logging


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





def 					__UpperCamelCase (     a__:			str      )					->      List[str]:




					"""simple docstring"""



					if "backbone.0.body" in orig_key:
										__lowercase        =      orig_key.replace('''backbone.0.body'''			,						'''backbone.conv_encoder.model'''      )
					if "transformer" in orig_key:
										__lowercase        =      orig_key.replace('''transformer.'''			,						''''''      )
					if "norm1" in orig_key:
										if "encoder" in orig_key:
															__lowercase        =      orig_key.replace('''norm1'''			,						'''self_attn_layer_norm'''      )
										else:
															__lowercase        =      orig_key.replace('''norm1'''			,						'''encoder_attn_layer_norm'''      )
					if "norm2" in orig_key:
										if "encoder" in orig_key:
															__lowercase        =      orig_key.replace('''norm2'''			,						'''final_layer_norm'''      )
										else:
															__lowercase        =      orig_key.replace('''norm2'''			,						'''self_attn_layer_norm'''      )
					if "norm3" in orig_key:
										__lowercase        =      orig_key.replace('''norm3'''			,						'''final_layer_norm'''      )
					if "linear1" in orig_key:
										__lowercase        =      orig_key.replace('''linear1'''			,						'''fc1'''      )
					if "linear2" in orig_key:
										__lowercase        =      orig_key.replace('''linear2'''			,						'''fc2'''      )
					if "query_embed" in orig_key:
										__lowercase        =      orig_key.replace('''query_embed'''			,						'''query_position_embeddings'''      )
					if "cross_attn" in orig_key:
										__lowercase        =      orig_key.replace('''cross_attn'''			,						'''encoder_attn'''      )

					return orig_key





def 					__UpperCamelCase (     a__:			Union[str, Any]      )					->      Tuple:




					"""simple docstring"""



					for i in range(6      ):
										# read in weights + bias of input projection layer of self-attention
										__lowercase        =      state_dict.pop(f'''decoder.layers.{i}.self_attn.in_proj_weight'''      )
										__lowercase        =      state_dict.pop(f'''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[:2_56, :]
										__lowercase        =      in_proj_bias[:2_56]
										__lowercase        =      in_proj_weight[2_56:5_12, :]
										__lowercase        =      in_proj_bias[2_56:5_12]
										__lowercase        =      in_proj_weight[-2_56:, :]
										__lowercase        =      in_proj_bias[-2_56:]





def 					__UpperCamelCase (     )					->      str:




					"""simple docstring"""



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

					return im





@torch.no_grad()
def 					__UpperCamelCase (     a__:			List[Any]			,						a__:			Union[str, Any]			,						a__:			Optional[int]			,						a__:			Optional[int]			,						a__:			List[Any]			,						a__:			Optional[Any]			,						a__:			str			,						)					->      List[str]:




					"""simple docstring"""



					__lowercase        =      DeformableDetrConfig()
					# set config attributes
					if single_scale:
										__lowercase        =      1
					__lowercase        =      dilation
					__lowercase        =      with_box_refine
					__lowercase        =      two_stage
					# set labels
					__lowercase        =      91
					__lowercase        =      '''huggingface/label-files'''
					__lowercase        =      '''coco-detection-id2label.json'''
					__lowercase        =      json.load(open(cached_download(hf_hub_url(a__			,						a__			,						repo_type='''dataset'''      )      )			,						'''r'''      )      )
					__lowercase        =      {int(a__      ): v for k, v in idalabel.items()}
					__lowercase        =      idalabel
					__lowercase        =      {v: k for k, v in idalabel.items()}

					# load image processor
					__lowercase        =      DeformableDetrImageProcessor(format='''coco_detection'''      )

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

					logger.info('''Converting model...'''      )

					# load original state dict
					__lowercase        =      torch.load(a__			,						map_location='''cpu'''      )['''model''']
					# rename keys
					for key in state_dict.copy().keys():
										__lowercase        =      state_dict.pop(a__      )
										__lowercase        =      val
					# query, key and value matrices need special treatment
					read_in_q_k_v(a__      )
					# important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
					__lowercase        =      '''model.'''
					for key in state_dict.copy().keys():
										if not key.startswith('''class_embed'''      ) and not key.startswith('''bbox_embed'''      ):
															__lowercase        =      state_dict.pop(a__      )
															__lowercase        =      val
    # finally, create HuggingFace model and load state dict
					__lowercase        =      DeformableDetrForObjectDetection(a__      )
					model.load_state_dict(a__      )
					model.eval()

					__lowercase        =      '''cuda''' if torch.cuda.is_available() else '''cpu'''
					model.to(a__      )
					# verify our conversion
					__lowercase        =      model(pixel_values.to(a__      )      )

					__lowercase        =      torch.tensor(
					    [[-9.66_45, -4.34_49, -5.87_05], [-9.70_35, -3.85_04, -5.07_24], [-10.56_34, -5.33_79, -7.51_16]]      )
					__lowercase        =      torch.tensor([[0.86_93, 0.22_89, 0.24_92], [0.31_50, 0.54_89, 0.58_45], [0.55_63, 0.75_80, 0.85_18]]      )

					if single_scale:
										__lowercase        =      torch.tensor(
										    [[-9.90_51, -4.25_41, -6.48_52], [-9.69_47, -4.08_54, -6.80_33], [-10.06_65, -5.84_70, -7.70_03]]      )
										__lowercase        =      torch.tensor([[0.72_92, 0.49_91, 0.55_32], [0.79_59, 0.24_26, 0.42_36], [0.75_82, 0.35_18, 0.44_51]]      )

					if single_scale and dilation:
										__lowercase        =      torch.tensor(
										    [[-8.96_52, -4.10_74, -5.66_35], [-9.05_96, -4.94_47, -6.60_75], [-10.11_78, -4.52_75, -6.26_71]]      )
										__lowercase        =      torch.tensor([[0.76_65, 0.41_30, 0.47_69], [0.83_64, 0.18_41, 0.33_91], [0.62_61, 0.38_95, 0.79_78]]      )

					if with_box_refine:
										__lowercase        =      torch.tensor(
										    [[-8.88_95, -5.41_87, -6.81_53], [-8.47_06, -6.16_68, -7.61_84], [-9.00_42, -5.53_59, -6.91_41]]      )
										__lowercase        =      torch.tensor([[0.78_28, 0.22_08, 0.43_23], [0.08_92, 0.59_96, 0.13_19], [0.55_24, 0.63_89, 0.89_14]]      )

					if with_box_refine and two_stage:
										__lowercase        =      torch.tensor(
										    [[-6.71_08, -4.32_13, -6.37_77], [-8.90_14, -6.17_99, -6.72_40], [-6.93_15, -4.47_35, -6.22_98]]      )
										__lowercase        =      torch.tensor([[0.25_83, 0.54_99, 0.46_83], [0.76_52, 0.90_68, 0.48_82], [0.54_90, 0.27_63, 0.05_64]]      )

					print('''Logits:'''			,						outputs.logits[0, :3, :3]      )

					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!'''      )

					# Save model and image processor
					logger.info(f'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'''      )
					Path(a__      ).mkdir(exist_ok=a__      )
					model.save_pretrained(a__      )
					image_processor.save_pretrained(a__      )

					# Push to hub
					if push_to_hub:
										__lowercase        =      '''deformable-detr'''
										model_name += "-single-scale" if single_scale else ""
										model_name += "-dc5" if dilation else ""
										model_name += "-with-box-refine" if with_box_refine else ""
										model_name += "-two-stage" if two_stage else ""
										print('''Pushing model to hub...'''      )
										model.push_to_hub(repo_path_or_name=a__			,						organization='''nielsr'''			,						commit_message='''Add model'''      )


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

							parser.add_argument(
							    """--checkpoint_path""",
							    type=str,
							    default="""/home/niels/checkpoints/deformable_detr/r50_deformable_detr-checkpoint.pth""",
							    help="""Path to Pytorch checkpoint (.pth file) you'd like to convert.""",
							)
							parser.add_argument("""--single_scale""", action="""store_true""", help="""Whether to set config.num_features_levels = 1.""")
							parser.add_argument("""--dilation""", action="""store_true""", help="""Whether to set config.dilation=True.""")
							parser.add_argument("""--with_box_refine""", action="""store_true""", help="""Whether to set config.with_box_refine=True.""")
							parser.add_argument("""--two_stage""", action="""store_true""", help="""Whether to set config.two_stage=True.""")
							parser.add_argument(
							    """--pytorch_dump_folder_path""",
							    default=None,
							    type=str,
							    required=True,
							    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__			:						Union[str, Any]							=			parser.parse_args()
							convert_deformable_detr_checkpoint(
							    args.checkpoint_path,
							    args.single_scale,
							    args.dilation,
							    args.with_box_refine,
							    args.two_stage,
							    args.pytorch_dump_folder_path,
							    args.push_to_hub,
							)


