





import argparse
import json
import pickle
from pathlib import Path

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

from transformers import MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, SwinConfig
from transformers.utils import logging


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






def lowerCamelCase__   (							snake_case_     :  str	)				->    Tuple:
   __snake_case								=       SwinConfig.from_pretrained(
       '''microsoft/swin-tiny-patch4-window7-224'''	,				out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']	)
   __snake_case								=       MaskFormerConfig(backbone_config=_UpperCamelCase	)

   __snake_case								=       '''huggingface/label-files'''
   if "ade20k-full" in model_name:
      # this should be ok
      __snake_case								=       847
      __snake_case								=       '''maskformer-ade20k-full-id2label.json'''
   elif "ade" in model_name:
      # this should be ok
      __snake_case								=       150
      __snake_case								=       '''ade20k-id2label.json'''
   elif "coco-stuff" in model_name:
      # this should be ok
      __snake_case								=       171
      __snake_case								=       '''maskformer-coco-stuff-id2label.json'''
   elif "coco" in model_name:
      # TODO
      __snake_case								=       133
      __snake_case								=       '''coco-panoptic-id2label.json'''
   elif "cityscapes" in model_name:
      # this should be ok
      __snake_case								=       19
      __snake_case								=       '''cityscapes-id2label.json'''
   elif "vistas" in model_name:
      # this should be ok
      __snake_case								=       65
      __snake_case								=       '''mapillary-vistas-id2label.json'''

   __snake_case								=       json.load(open(hf_hub_download(_UpperCamelCase	,				_UpperCamelCase	,				repo_type='''dataset'''	)	,				'''r'''	)	)
   __snake_case								=       {int(_UpperCamelCase	): v for k, v in idalabel.items()}

   return config






def lowerCamelCase__   (							snake_case_     :  int	)				->    Tuple:
   __snake_case								=       []
   # stem
   # fmt: off
   rename_keys.append(('''backbone.patch_embed.proj.weight''', '''model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.weight''')	)
   rename_keys.append(('''backbone.patch_embed.proj.bias''', '''model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.bias''')	)
   rename_keys.append(('''backbone.patch_embed.norm.weight''', '''model.pixel_level_module.encoder.model.embeddings.norm.weight''')	)
   rename_keys.append(('''backbone.patch_embed.norm.bias''', '''model.pixel_level_module.encoder.model.embeddings.norm.bias''')	)
   # stages
   for i in range(len(config.backbone_config.depths	)	):
      for j in range(config.backbone_config.depths[i]	):
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.norm1.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.norm1.bias""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.attn.relative_position_bias_table""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.attn.relative_position_index""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.attn.proj.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.attn.proj.bias""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.norm2.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.norm2.bias""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.mlp.fc1.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.mlp.fc1.bias""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.mlp.fc2.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.blocks.{j}.mlp.fc2.bias""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.bias""")	)

      if i < 3:
         rename_keys.append((f"""backbone.layers.{i}.downsample.reduction.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.reduction.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.downsample.norm.weight""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.weight""")	)
         rename_keys.append((f"""backbone.layers.{i}.downsample.norm.bias""", f"""model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.bias""")	)
      rename_keys.append((f"""backbone.norm{i}.weight""", f"""model.pixel_level_module.encoder.hidden_states_norms.{i}.weight""")	)
      rename_keys.append((f"""backbone.norm{i}.bias""", f"""model.pixel_level_module.encoder.hidden_states_norms.{i}.bias""")	)

   # FPN
   rename_keys.append(('''sem_seg_head.layer_4.weight''', '''model.pixel_level_module.decoder.fpn.stem.0.weight''')	)
   rename_keys.append(('''sem_seg_head.layer_4.norm.weight''', '''model.pixel_level_module.decoder.fpn.stem.1.weight''')	)
   rename_keys.append(('''sem_seg_head.layer_4.norm.bias''', '''model.pixel_level_module.decoder.fpn.stem.1.bias''')	)
   for source_index, target_index in zip(range(3	,				0	,				-1	)	,				range(0	,				3	)	):
      rename_keys.append((f"""sem_seg_head.adapter_{source_index}.weight""", f"""model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight""")	)
      rename_keys.append((f"""sem_seg_head.adapter_{source_index}.norm.weight""", f"""model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight""")	)
      rename_keys.append((f"""sem_seg_head.adapter_{source_index}.norm.bias""", f"""model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias""")	)
      rename_keys.append((f"""sem_seg_head.layer_{source_index}.weight""", f"""model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight""")	)
      rename_keys.append((f"""sem_seg_head.layer_{source_index}.norm.weight""", f"""model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight""")	)
      rename_keys.append((f"""sem_seg_head.layer_{source_index}.norm.bias""", f"""model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias""")	)
   rename_keys.append(('''sem_seg_head.mask_features.weight''', '''model.pixel_level_module.decoder.mask_projection.weight''')	)
   rename_keys.append(('''sem_seg_head.mask_features.bias''', '''model.pixel_level_module.decoder.mask_projection.bias''')	)

   # Transformer decoder
   for idx in range(config.decoder_config.decoder_layers	):
      # self-attention out projection
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight""", f"""model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias""", f"""model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias""")	)
      # cross-attention out projection
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight""", f"""model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias""", f"""model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias""")	)
      # MLP 1
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight""", f"""model.transformer_module.decoder.layers.{idx}.fc1.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias""", f"""model.transformer_module.decoder.layers.{idx}.fc1.bias""")	)
      # MLP 2
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight""", f"""model.transformer_module.decoder.layers.{idx}.fc2.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias""", f"""model.transformer_module.decoder.layers.{idx}.fc2.bias""")	)
      # layernorm 1 (self-attention layernorm)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight""", f"""model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias""", f"""model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias""")	)
      # layernorm 2 (cross-attention layernorm)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight""", f"""model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias""", f"""model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias""")	)
      # layernorm 3 (final layernorm)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight""", f"""model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias""", f"""model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias""")	)

   rename_keys.append(('''sem_seg_head.predictor.transformer.decoder.norm.weight''', '''model.transformer_module.decoder.layernorm.weight''')	)
   rename_keys.append(('''sem_seg_head.predictor.transformer.decoder.norm.bias''', '''model.transformer_module.decoder.layernorm.bias''')	)

   # heads on top
   rename_keys.append(('''sem_seg_head.predictor.query_embed.weight''', '''model.transformer_module.queries_embedder.weight''')	)

   rename_keys.append(('''sem_seg_head.predictor.input_proj.weight''', '''model.transformer_module.input_projection.weight''')	)
   rename_keys.append(('''sem_seg_head.predictor.input_proj.bias''', '''model.transformer_module.input_projection.bias''')	)

   rename_keys.append(('''sem_seg_head.predictor.class_embed.weight''', '''class_predictor.weight''')	)
   rename_keys.append(('''sem_seg_head.predictor.class_embed.bias''', '''class_predictor.bias''')	)

   for i in range(3	):
      rename_keys.append((f"""sem_seg_head.predictor.mask_embed.layers.{i}.weight""", f"""mask_embedder.{i}.0.weight""")	)
      rename_keys.append((f"""sem_seg_head.predictor.mask_embed.layers.{i}.bias""", f"""mask_embedder.{i}.0.bias""")	)
   # fmt: on

   return rename_keys






def lowerCamelCase__   (							snake_case_     :  Union[str, Any]	,				snake_case_     :  List[Any]	,				snake_case_     :  Optional[int]	)				->    Union[str, Any]:
   __snake_case								=       dct.pop(_UpperCamelCase	)
   __snake_case								=       val






def lowerCamelCase__   (							snake_case_     :  Optional[int]	,				snake_case_     :  int	)				->    int:
   __snake_case								=       [int(backbone_config.embed_dim * 2**i	) for i in range(len(backbone_config.depths	)	)]
   for i in range(len(backbone_config.depths	)	):
      __snake_case								=       num_features[i]
      for j in range(backbone_config.depths[i]	):
         # fmt: off
         # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
         __snake_case								=       state_dict.pop(f"""backbone.layers.{i}.blocks.{j}.attn.qkv.weight"""	)
         __snake_case								=       state_dict.pop(f"""backbone.layers.{i}.blocks.{j}.attn.qkv.bias"""	)
         # next, add query, keys and values (in that order) to the state dict
         __snake_case								=       in_proj_weight[:dim, :]
         __snake_case								=       in_proj_bias[: dim]
         __snake_case								=       in_proj_weight[
             dim : dim * 2, :
         ]
         __snake_case								=       in_proj_bias[
             dim : dim * 2
         ]
         __snake_case								=       in_proj_weight[
             -dim :, :
         ]
         __snake_case								=       in_proj_bias[-dim :]
         # fmt: on






def lowerCamelCase__   (							snake_case_     :  str	,				snake_case_     :  Optional[Any]	)				->    str:
   __snake_case								=       config.decoder_config.hidden_size
   for idx in range(config.decoder_config.decoder_layers	):
      # read in weights + bias of self-attention input projection layer (in the original implementation, this is a single matrix + bias)
      __snake_case								=       state_dict.pop(f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight"""	)
      __snake_case								=       state_dict.pop(f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.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[:config.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 :]
      # read in weights + bias of cross-attention input projection layer (in the original implementation, this is a single matrix + bias)
      __snake_case								=       state_dict.pop(f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight"""	)
      __snake_case								=       state_dict.pop(f"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_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[:config.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 :]
   # fmt: on






def lowerCamelCase__   (							)				->    torch.Tensor:
   __snake_case								=       '''http://images.cocodataset.org/val2017/000000039769.jpg'''
   __snake_case								=       Image.open(requests.get(_UpperCamelCase	,				stream=_UpperCamelCase	).raw	)
   return im






@torch.no_grad()
def lowerCamelCase__   (							snake_case_     :  str	,				snake_case_     :  str	,				snake_case_     :  str	,				snake_case_     :  bool = False	)				->    Dict:
   __snake_case								=       get_maskformer_config(_UpperCamelCase	)

   # load original state_dict
   with open(_UpperCamelCase	,				'''rb'''	) as f:
      __snake_case								=       pickle.load(_UpperCamelCase	)
   __snake_case								=       data['''model''']

   # for name, param in state_dict.items():
   #     print(name, param.shape)

   # rename keys
   __snake_case								=       create_rename_keys(_UpperCamelCase	)
   for src, dest in rename_keys:
      rename_key(_UpperCamelCase	,				_UpperCamelCase	,				_UpperCamelCase	)
   read_in_swin_q_k_v(_UpperCamelCase	,				config.backbone_config	)
   read_in_decoder_q_k_v(_UpperCamelCase	,				_UpperCamelCase	)

   # update to torch tensors
   for key, value in state_dict.items():
      __snake_case								=       torch.from_numpy(_UpperCamelCase	)

   # load 🤗 model
   __snake_case								=       MaskFormerForInstanceSegmentation(_UpperCamelCase	)
   model.eval()

   for name, param in model.named_parameters():
      print(_UpperCamelCase	,				param.shape	)

   __snake_case	,     __snake_case								=       model.load_state_dict(_UpperCamelCase	,				strict=_UpperCamelCase	)
   assert missing_keys == [
       "model.pixel_level_module.encoder.model.layernorm.weight",
       "model.pixel_level_module.encoder.model.layernorm.bias",
   ]
   assert len(_UpperCamelCase	) == 0, f"""Unexpected keys: {unexpected_keys}"""

   # verify results
   __snake_case								=       prepare_img()
   if "vistas" in model_name:
      __snake_case								=       65
   elif "cityscapes" in model_name:
      __snake_case								=       6_5535
   else:
      __snake_case								=       255
   __snake_case								=       True if '''ade''' in model_name else False
   __snake_case								=       MaskFormerImageProcessor(ignore_index=_UpperCamelCase	,				reduce_labels=_UpperCamelCase	)

   __snake_case								=       image_processor(_UpperCamelCase	,				return_tensors='''pt'''	)

   __snake_case								=       model(**_UpperCamelCase	)

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

   if model_name == "maskformer-swin-tiny-ade":
      __snake_case								=       torch.tensor(
          [[3.6_353, -4.4_770, -2.6_065], [0.5_081, -4.2_394, -3.5_343], [2.1_909, -5.0_353, -1.9_323]]	)
   assert torch.allclose(outputs.class_queries_logits[0, :3, :3]	,				_UpperCamelCase	,				atol=1e-4	)
   print('''Looks ok!'''	)

   if pytorch_dump_folder_path is not None:
      print(f"""Saving model and image processor to {pytorch_dump_folder_path}"""	)
      Path(_UpperCamelCase	).mkdir(exist_ok=_UpperCamelCase	)
      model.save_pretrained(_UpperCamelCase	)
      image_processor.save_pretrained(_UpperCamelCase	)

   if push_to_hub:
      print('''Pushing model and image processor to the hub...'''	)
      model.push_to_hub(f"""nielsr/{model_name}"""	)
      image_processor.push_to_hub(f"""nielsr/{model_name}"""	)


if __name__ == "__main__":
    snake_case_           =      argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--model_name',
        default='maskformer-swin-tiny-ade',
        type=str,
        help=('Name of the MaskFormer model you\'d like to convert',),
    )
    parser.add_argument(
        '--checkpoint_path',
        default='/Users/nielsrogge/Documents/MaskFormer_checkpoints/MaskFormer-Swin-tiny-ADE20k/model.pkl',
        type=str,
        help='Path to the original state dict (.pth file).',
    )
    parser.add_argument(
        '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
    )
    parser.add_argument(
        '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
    )

    snake_case_           =      parser.parse_args()
    convert_maskformer_checkpoint(
        args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
    )








# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
    import torch

class 		SCREAMING_SNAKE_CASE__    (	TensorFormatter[Mapping, 'torch.Tensor', Mapping]  ):


      def __init__(self	:				Optional[int]    ,      a__	:				Tuple=None    ,      **a__	:				Optional[int]				):





         """simple docstring"""



         super().__init__(features=a__				)
         __snake_case								=       torch_tensor_kwargs
         import torch  # noqa import torch at initialization


      def        a			(self	:				Union[str, Any]    ,      a__	:				Union[str, Any]				):





         """simple docstring"""



         import torch

         if isinstance(a__    ,      a__				) and column:
            if all(
                isinstance(a__    ,      torch.Tensor				) and x.shape == column[0].shape and x.dtype == column[0].dtype
                for x in column				):
               return torch.stack(a__				)
         return column


      def        a			(self	:				Optional[Any]    ,      a__	:				str				):





         """simple docstring"""



         import torch

         if isinstance(a__    ,      (str, bytes, type(a__				))				):
            return value
         elif isinstance(a__    ,      (np.character, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.character				):
            return value.tolist()

         __snake_case								=       {}

         if isinstance(a__    ,      (np.number, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.integer				):
            __snake_case								=       {'''dtype''': torch.intaa}
         elif isinstance(a__    ,      (np.number, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.floating				):
            __snake_case								=       {'''dtype''': torch.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
            import PIL.Image

            if isinstance(a__    ,      PIL.Image.Image				):
               __snake_case								=       np.asarray(a__				)
         return torch.tensor(a__    ,      **{**default_dtype, **self.torch_tensor_kwargs}				)


      def        a			(self	:				Optional[int]    ,      a__	:				str				):





         """simple docstring"""



         import torch

         # support for torch, tf, jax etc.
         if hasattr(a__    ,      '''__array__'''				) and not isinstance(a__    ,      torch.Tensor				):
            __snake_case								=       data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(a__    ,      np.ndarray				):
            if data_struct.dtype == object:  # torch tensors cannot be instantied from an array of objects
               return self._consolidate([self.recursive_tensorize(a__				) for substruct in data_struct]				)
         elif isinstance(a__    ,      (list, tuple)				):
            return self._consolidate([self.recursive_tensorize(a__				) for substruct in data_struct]				)
         return self._tensorize(a__				)


      def        a			(self	:				Optional[Any]    ,      a__	:				dict				):





         """simple docstring"""



         return map_nested(self._recursive_tensorize    ,      a__    ,      map_list=a__				)


      def        a			(self	:				Optional[int]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_row(a__				)
         __snake_case								=       self.python_features_decoder.decode_row(a__				)
         return self.recursive_tensorize(a__				)


      def        a			(self	:				List[Any]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_column(a__				)
         __snake_case								=       self.python_features_decoder.decode_column(a__    ,      pa_table.column_names[0]				)
         __snake_case								=       self.recursive_tensorize(a__				)
         __snake_case								=       self._consolidate(a__				)
         return column


      def        a			(self	:				Optional[int]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_batch(a__				)
         __snake_case								=       self.python_features_decoder.decode_batch(a__				)
         __snake_case								=       self.recursive_tensorize(a__				)
         for column_name in batch:
            __snake_case								=       self._consolidate(batch[column_name]				)
         return batch

