





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 DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
from transformers.utils import logging


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




def    __UpperCAmelCase(      lowercase_      ):


       if "resnet-50" in model_name:
              _lowerCamelCase							:     List[Any]	=		ResNetConfig.from_pretrained('''microsoft/resnet-50'''      )
       elif "resnet-101" in model_name:
              _lowerCamelCase							:     Optional[int]	=		ResNetConfig.from_pretrained('''microsoft/resnet-101'''      )
       else:
              raise ValueError('''Model name should include either resnet50 or resnet101'''      )

       _lowerCamelCase							:     str	=		DetrConfig(use_timm_backbone=__lowerCAmelCase ,    backbone_config=__lowerCAmelCase      )

       # set label attributes
       _lowerCamelCase							:     List[str]	=		'''panoptic''' in model_name
       if is_panoptic:
              _lowerCamelCase							:     Tuple	=		2_50
       else:
              _lowerCamelCase							:     Tuple	=		91
              _lowerCamelCase							:     List[str]	=		'''huggingface/label-files'''
              _lowerCamelCase							:     Dict	=		'''coco-detection-id2label.json'''
              _lowerCamelCase							:     Optional[Any]	=		json.load(open(hf_hub_download(__lowerCAmelCase ,    __lowerCAmelCase ,    repo_type='''dataset'''      ) ,    '''r'''      )      )
              _lowerCamelCase							:     Dict	=		{int(__lowerCAmelCase      ): v for k, v in idalabel.items()}
              _lowerCamelCase							:     List[Any]	=		idalabel
              _lowerCamelCase							:     List[str]	=		{v: k for k, v in idalabel.items()}

       return config, is_panoptic




def    __UpperCAmelCase(      lowercase_      ):
       _lowerCamelCase							:     List[Any]	=		[]

       # stem
       # fmt: off
       rename_keys.append(('''backbone.0.body.conv1.weight''', '''backbone.conv_encoder.model.embedder.embedder.convolution.weight''')      )
       rename_keys.append(('''backbone.0.body.bn1.weight''', '''backbone.conv_encoder.model.embedder.embedder.normalization.weight''')      )
       rename_keys.append(('''backbone.0.body.bn1.bias''', '''backbone.conv_encoder.model.embedder.embedder.normalization.bias''')      )
       rename_keys.append(('''backbone.0.body.bn1.running_mean''', '''backbone.conv_encoder.model.embedder.embedder.normalization.running_mean''')      )
       rename_keys.append(('''backbone.0.body.bn1.running_var''', '''backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.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"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var""",
                                )      )
              # fmt: on

       for i in range(config.encoder_layers      ):
              # 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'''),
           ]      )

       return rename_keys




def    __UpperCAmelCase(      lowercase_ ,    lowercase_ ,    lowercase_      ):
       _lowerCamelCase							:     Any	=		state_dict.pop(__lowerCAmelCase      )
       _lowerCamelCase							:     Any	=		val




def    __UpperCAmelCase(      lowercase_ ,    lowercase_=False      ):
       _lowerCamelCase							:     Tuple	=		''''''
       if is_panoptic:
              _lowerCamelCase							:     List[str]	=		'''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)
              _lowerCamelCase							:     Union[str, Any]	=		state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""      )
              _lowerCamelCase							:     Tuple	=		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
              _lowerCamelCase							:     Union[str, Any]	=		in_proj_weight[:2_56, :]
              _lowerCamelCase							:     Dict	=		in_proj_bias[:2_56]
              _lowerCamelCase							:     Dict	=		in_proj_weight[2_56:5_12, :]
              _lowerCamelCase							:     int	=		in_proj_bias[2_56:5_12]
              _lowerCamelCase							:     Union[str, Any]	=		in_proj_weight[-2_56:, :]
              _lowerCamelCase							:     Tuple	=		in_proj_bias[-2_56:]
       # 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
              _lowerCamelCase							:     Any	=		state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""      )
              _lowerCamelCase							:     str	=		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
              _lowerCamelCase							:     Optional[int]	=		in_proj_weight[:2_56, :]
              _lowerCamelCase							:     str	=		in_proj_bias[:2_56]
              _lowerCamelCase							:     Optional[int]	=		in_proj_weight[2_56:5_12, :]
              _lowerCamelCase							:     Optional[int]	=		in_proj_bias[2_56:5_12]
              _lowerCamelCase							:     Dict	=		in_proj_weight[-2_56:, :]
              _lowerCamelCase							:     Union[str, Any]	=		in_proj_bias[-2_56:]
              # read in weights + bias of input projection layer of cross-attention
              _lowerCamelCase							:     List[Any]	=		state_dict.pop(
                  F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""      )
              _lowerCamelCase							:     str	=		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
              _lowerCamelCase							:     List[str]	=		in_proj_weight_cross_attn[:2_56, :]
              _lowerCamelCase							:     Tuple	=		in_proj_bias_cross_attn[:2_56]
              _lowerCamelCase							:     Any	=		in_proj_weight_cross_attn[2_56:5_12, :]
              _lowerCamelCase							:     Dict	=		in_proj_bias_cross_attn[2_56:5_12]
              _lowerCamelCase							:     str	=		in_proj_weight_cross_attn[-2_56:, :]
              _lowerCamelCase							:     List[Any]	=		in_proj_bias_cross_attn[-2_56:]




def    __UpperCAmelCase(      ):
       _lowerCamelCase							:     str	=		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
       _lowerCamelCase							:     Dict	=		Image.open(requests.get(__lowerCAmelCase ,    stream=__lowerCAmelCase      ).raw      )

       return im




@torch.no_grad()
def    __UpperCAmelCase(      lowercase_ ,    lowercase_=None ,    lowercase_=False      ):
       _lowerCamelCase,  _lowerCamelCase							:     str	=		get_detr_config(__lowerCAmelCase      )

       # load original model from torch hub
       _lowerCamelCase							:     Optional[Any]	=		{
           '''detr-resnet-50''': '''detr_resnet50''',
           '''detr-resnet-101''': '''detr_resnet101''',
       }
       logger.info(F"""Converting model {model_name}..."""      )
       _lowerCamelCase							:     Any	=		torch.hub.load('''facebookresearch/detr''' ,    model_name_to_original_name[model_name] ,    pretrained=__lowerCAmelCase      ).eval()
       _lowerCamelCase							:     Union[str, Any]	=		detr.state_dict()
       # rename keys
       for src, dest in create_rename_keys(__lowerCAmelCase      ):
              if is_panoptic:
                     _lowerCamelCase							:     str	=		'''detr.''' + src
              rename_key(__lowerCAmelCase ,    __lowerCAmelCase ,    __lowerCAmelCase      )
       # query, key and value matrices need special treatment
       read_in_q_k_v(__lowerCAmelCase ,    is_panoptic=__lowerCAmelCase      )
       # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
       _lowerCamelCase							:     str	=		'''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'''      )
                     ):
                            _lowerCamelCase							:     Optional[int]	=		state_dict.pop(__lowerCAmelCase      )
                            _lowerCamelCase							:     str	=		val
                     elif "class_labels_classifier" in key or "bbox_predictor" in key:
                            _lowerCamelCase							:     List[str]	=		state_dict.pop(__lowerCAmelCase      )
                            _lowerCamelCase							:     List[Any]	=		val
                     elif key.startswith('''bbox_attention'''      ) or key.startswith('''mask_head'''      ):
                            continue
                     else:
                            _lowerCamelCase							:     Union[str, Any]	=		state_dict.pop(__lowerCAmelCase      )
                            _lowerCamelCase							:     List[Any]	=		val
              else:
                     if not key.startswith('''class_labels_classifier'''      ) and not key.startswith('''bbox_predictor'''      ):
                            _lowerCamelCase							:     Optional[int]	=		state_dict.pop(__lowerCAmelCase      )
                            _lowerCamelCase							:     str	=		val

              # finally, create HuggingFace model and load state dict
       _lowerCamelCase							:     Optional[int]	=		DetrForSegmentation(__lowerCAmelCase      ) if is_panoptic else DetrForObjectDetection(__lowerCAmelCase      )
       model.load_state_dict(__lowerCAmelCase      )
       model.eval()

       # verify our conversion on an image
       _lowerCamelCase							:     List[Any]	=		'''coco_panoptic''' if is_panoptic else '''coco_detection'''
       _lowerCamelCase							:     str	=		DetrImageProcessor(format=__lowerCAmelCase      )

       _lowerCamelCase							:     List[str]	=		processor(images=prepare_img() ,    return_tensors='''pt'''      )
       _lowerCamelCase							:     List[Any]	=		encoding['''pixel_values''']

       _lowerCamelCase							:     Union[str, Any]	=		detr(__lowerCAmelCase      )
       _lowerCamelCase							:     Tuple	=		model(__lowerCAmelCase      )

       assert torch.allclose(outputs.logits ,    original_outputs['''pred_logits'''] ,    atol=1e-3      )
       assert torch.allclose(outputs.pred_boxes ,    original_outputs['''pred_boxes'''] ,    atol=1e-3      )
       if is_panoptic:
              assert torch.allclose(outputs.pred_masks ,    original_outputs['''pred_masks'''] ,    atol=1e-4      )
       print('''Looks ok!'''      )

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

       if push_to_hub:
              # Upload model and image processor to the hub
              logger.info('''Uploading PyTorch model and image processor to the hub...'''      )
              model.push_to_hub(F"""nielsr/{model_name}"""      )
              processor.push_to_hub(F"""nielsr/{model_name}"""      )


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

       parser.add_argument(
           '--model_name',
           default='detr-resnet-50',
           type=str,
           choices=['detr-resnet-50', 'detr-resnet-101'],
           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.'
       )
       parser.add_argument('--push_to_hub', action='store_true', help='Whether to push the model to the hub or not.')
       _lowerCamelCase			       =      parser.parse_args()
       convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch




def    __UpperCAmelCase(      lowercase_      ):  # picklable for multiprocessing
       return x.sum()




def    __UpperCAmelCase(      lowercase_      ):  # picklable for multiprocessing
       return i + 1







@dataclass
class 			__A       :




      """simple docstring"""







      UpperCAmelCase__						=      42
      UpperCAmelCase__						=      42







class 			__A       (      lowerCamelCase__  ):




      """simple docstring"""







      def 				__snake_case			(		self):



             """simple docstring"""





             _lowerCamelCase							:     Tuple	=		{}
             _lowerCamelCase							:     Dict	=		[]
             _lowerCamelCase							:     Optional[Any]	=		1
             _lowerCamelCase							:     Optional[int]	=		[1, 2]
             _lowerCamelCase							:     str	=		{'''a''': 1, '''b''': 2}
             _lowerCamelCase							:     Dict	=		{'''a''': [1, 2], '''b''': [3, 4]}
             _lowerCamelCase							:     Any	=		{'''a''': {'''1''': 1}, '''b''': 2}
             _lowerCamelCase							:     Optional[Any]	=		{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
             _lowerCamelCase							:     str	=		{}
             _lowerCamelCase							:     int	=		[]
             _lowerCamelCase							:     str	=		2
             _lowerCamelCase							:     int	=		[2, 3]
             _lowerCamelCase							:     str	=		{'''a''': 2, '''b''': 3}
             _lowerCamelCase							:     Tuple	=		{'''a''': [2, 3], '''b''': [4, 5]}
             _lowerCamelCase							:     List[str]	=		{'''a''': {'''1''': 2}, '''b''': 3}
             _lowerCamelCase							:     str	=		{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__)      ,							a__)

             _lowerCamelCase							:     Dict	=		2
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(map_nested(a__      ,							a__      ,							num_proc=a__)      ,							a__)

             _lowerCamelCase							:     Any	=		{'''a''': np.eye(2), '''b''': np.zeros(3), '''c''': np.ones(2)}
             _lowerCamelCase							:     Optional[int]	=		{'''a''': 2, '''b''': 0, '''c''': 2}
             _lowerCamelCase							:     Optional[int]	=		{
                 '''a''': np.eye(2).astype(a__),
                 '''b''': np.zeros(3).astype(a__),
                 '''c''': np.ones(2).astype(a__),
             }
             self.assertEqual(map_nested(a__      ,							a__      ,							map_numpy=a__)      ,							a__)
             self.assertEqual(
                 {k: v.tolist() for k, v in map_nested(a__      ,							a__      ,							map_numpy=a__).items()}      ,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,							)
             self.assertEqual(map_nested(a__      ,							a__      ,							map_numpy=a__      ,							num_proc=a__)      ,							a__)
             self.assertEqual(
                 {k: v.tolist() for k, v in map_nested(a__      ,							a__      ,							map_numpy=a__      ,							num_proc=a__).items()}      ,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,							)
             with self.assertRaises(a__):  # can't pickle a local lambda
                    map_nested(lambda a__: x + 1      ,							a__      ,							num_proc=a__)


      def 				__snake_case			(		self):



             """simple docstring"""





             _lowerCamelCase							:     Dict	=		{'''a''': 1, '''b''': 2}
             _lowerCamelCase							:     Optional[int]	=		{'''a''': 3, '''b''': 4}
             _lowerCamelCase							:     int	=		{'''a''': 5, '''b''': 6}
             _lowerCamelCase							:     Optional[int]	=		sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))])
             self.assertEqual(sorted(zip_dict(a__      ,							a__      ,							a__))      ,							a__)



      def 				__snake_case			(		self):



             """simple docstring"""





             class 			__A       :




                   """simple docstring"""







                   UpperCAmelCase__						=      """bar"""
             _lowerCamelCase							:     Any	=		Foo()
             self.assertEqual(foo.my_attr      ,							'''bar''')
             with temporary_assignment(a__      ,							'''my_attr'''      ,							'''BAR'''):
                    self.assertEqual(foo.my_attr      ,							'''BAR''')
             self.assertEqual(foo.my_attr      ,							'''bar''')




@pytest.mark.parametrize(
    '''iterable_length, num_proc, expected_num_proc''' ,    [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ] ,    )
def    __UpperCAmelCase(      lowercase_ ,    lowercase_ ,    lowercase_      ):
       with patch('''datasets.utils.py_utils._single_map_nested'''      ) as mock_single_map_nested, patch(
           '''datasets.parallel.parallel.Pool'''      ) as mock_multiprocessing_pool:
              _lowerCamelCase							:     Union[str, Any]	=		{F"""{i}""": i for i in range(lowercase_      )}
              _lowerCamelCase							:     List[str]	=		map_nested(lambda lowercase_      : x + 10 ,    lowercase_ ,    num_proc=lowercase_ ,    parallel_min_length=16      )
              if expected_num_proc == 1:
                     assert mock_single_map_nested.called
                     assert not mock_multiprocessing_pool.called
              else:
                     assert not mock_single_map_nested.called
                     assert mock_multiprocessing_pool.called
                     assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc







class 			__A       (      lowerCamelCase__  ):




      """simple docstring"""







      @require_tf
      def 				__snake_case			(		self):



             """simple docstring"""





             import tensorflow as tf
             from tensorflow.keras import layers

             _lowerCamelCase							:     int	=		layers.Dense(2)

             def gen_random_output():
                    _lowerCamelCase							:     Union[str, Any]	=		tf.random.uniform((1, 3))
                    return model(a__).numpy()

             with temp_seed(42      ,							set_tensorflow=a__):
                    _lowerCamelCase							:     List[str]	=		gen_random_output()
             with temp_seed(42      ,							set_tensorflow=a__):
                    _lowerCamelCase							:     Any	=		gen_random_output()
             _lowerCamelCase							:     str	=		gen_random_output()

             np.testing.assert_equal(a__      ,							a__)
             self.assertGreater(np.abs(outa - outa).sum()      ,							0)


      @require_torch
      def 				__snake_case			(		self):



             """simple docstring"""





             import torch

             def gen_random_output():
                    _lowerCamelCase							:     Union[str, Any]	=		torch.nn.Linear(3      ,							2)
                    _lowerCamelCase							:     Dict	=		torch.rand(1      ,							3)
                    return model(a__).detach().numpy()

             with temp_seed(42      ,							set_pytorch=a__):
                    _lowerCamelCase							:     Any	=		gen_random_output()
             with temp_seed(42      ,							set_pytorch=a__):
                    _lowerCamelCase							:     Optional[int]	=		gen_random_output()
             _lowerCamelCase							:     Union[str, Any]	=		gen_random_output()

             np.testing.assert_equal(a__      ,							a__)
             self.assertGreater(np.abs(outa - outa).sum()      ,							0)



      def 				__snake_case			(		self):



             """simple docstring"""





             def gen_random_output():
                    return np.random.rand(1      ,							3)

             with temp_seed(42):
                    _lowerCamelCase							:     Union[str, Any]	=		gen_random_output()
             with temp_seed(42):
                    _lowerCamelCase							:     List[str]	=		gen_random_output()
             _lowerCamelCase							:     str	=		gen_random_output()

             np.testing.assert_equal(a__      ,							a__)
             self.assertGreater(np.abs(outa - outa).sum()      ,							0)




@pytest.mark.parametrize('''input_data''' ,    [{}]      )
def    __UpperCAmelCase(      lowercase_      ):
       _lowerCamelCase							:     List[Any]	=		NestedDataStructure(lowercase_      ).data
       assert output_data == input_data




@pytest.mark.parametrize(
    '''data, expected_output''' ,    [
        ({}, []),
        ([], []),
        ('''foo''', ['''foo''']),
        (['''foo''', '''bar'''], ['''foo''', '''bar''']),
        ([['''foo''', '''bar''']], ['''foo''', '''bar''']),
        ([[['''foo'''], ['''bar''']]], ['''foo''', '''bar''']),
        ([[['''foo'''], '''bar''']], ['''foo''', '''bar''']),
        ({'''a''': 1, '''b''': 2}, [1, 2]),
        ({'''a''': [1, 2], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, [4]]}, [1, 2, 3, 4]),
        ({'''a''': {'''1''': 1}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': [2]}, [1, 2]),
    ] ,    )
def    __UpperCAmelCase(      lowercase_ ,    lowercase_      ):
       _lowerCamelCase							:     int	=		NestedDataStructure(lowercase_      ).flatten()
       assert output == expected_output




def    __UpperCAmelCase(      ):
       _lowerCamelCase							:     Any	=		A(x=1 ,    y='''foobar'''      )
       _lowerCamelCase							:     Union[str, Any]	=		{'''x''': 1, '''y''': '''foobar'''}
       assert asdict(lowercase_      ) == expected_output

       _lowerCamelCase							:     Optional[int]	=		{'''a''': {'''b''': A(x=10 ,    y='''foo'''      )}, '''c''': [A(x=20 ,    y='''bar'''      )]}
       _lowerCamelCase							:     Union[str, Any]	=		{'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
       assert asdict(lowercase_      ) == expected_output

       with pytest.raises(lowercase_      ):
              asdict([1, A(x=10 ,    y='''foo'''      )]      )




def    __UpperCAmelCase(      lowercase_      ):
       return text.split()




def    __UpperCAmelCase(      lowercase_      ):
       yield (time.time(), content)
       time.sleep(2      )
       yield (time.time(), content)




def    __UpperCAmelCase(      ):
       with Pool(2      ) as pool:
              _lowerCamelCase							:     Tuple	=		list(iflatmap_unordered(lowercase_ ,    _split_text ,    kwargs_iterable=[{'''text''': '''hello there'''}] * 10      )      )
              assert out.count('''hello'''      ) == 10
              assert out.count('''there'''      ) == 10
              assert len(lowercase_      ) == 20

       # check multiprocess from pathos (uses dill for pickling)
       with multiprocess.Pool(2      ) as pool:
              _lowerCamelCase							:     Dict	=		list(iflatmap_unordered(lowercase_ ,    _split_text ,    kwargs_iterable=[{'''text''': '''hello there'''}] * 10      )      )
              assert out.count('''hello'''      ) == 10
              assert out.count('''there'''      ) == 10
              assert len(lowercase_      ) == 20

       # check that we get items as fast as possible
       with Pool(2      ) as pool:
              _lowerCamelCase							:     str	=		[]
              for yield_time, content in iflatmap_unordered(
                  lowercase_ ,    _aseconds_generator_of_aitems_with_timing ,    kwargs_iterable=[{'''content''': '''a'''}, {'''content''': '''b'''}]      ):
                     assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
                     out.append(lowercase_      )
              assert out.count('''a'''      ) == 2
              assert out.count('''b'''      ) == 2
              assert len(lowercase_      ) == 4

