






'''simple docstring'''





import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)






def    lowerCamelCase_     (    lowercase__				,     lowercase__):
 lowerCamelCase__  =     XCLIPTextConfig()

 # derive patch size from model name
 lowerCamelCase__  =     model_name.find("patch")
 lowerCamelCase__  =     int(model_name[start_idx + len("patch") : start_idx + len("patch") + 2])
 lowerCamelCase__  =     XCLIPVisionConfig(patch_size=lowercase__				,     num_frames=lowercase__)

 if "large" in model_name:
  lowerCamelCase__  =     768
  lowerCamelCase__  =     3072
  lowerCamelCase__  =     12

  lowerCamelCase__  =     1024
  lowerCamelCase__  =     4096
  lowerCamelCase__  =     16
  lowerCamelCase__  =     24
  lowerCamelCase__  =     768
  lowerCamelCase__  =     3072

 if model_name == "xclip-large-patch14-16-frames":
  lowerCamelCase__  =     336

 lowerCamelCase__  =     XCLIPConfig.from_text_vision_configs(lowercase__				,     lowercase__)

 if "large" in model_name:
  lowerCamelCase__  =     768

 return config






def    lowerCamelCase_     (    lowercase__):
 # text encoder
 if name == "token_embedding.weight":
  lowerCamelCase__  =     name.replace("token_embedding.weight"				,     "text_model.embeddings.token_embedding.weight")
 if name == "positional_embedding":
  lowerCamelCase__  =     name.replace("positional_embedding"				,     "text_model.embeddings.position_embedding.weight")
 if "ln_1" in name:
  lowerCamelCase__  =     name.replace("ln_1"				,     "layer_norm1")
 if "ln_2" in name:
  lowerCamelCase__  =     name.replace("ln_2"				,     "layer_norm2")
 if "c_fc" in name:
  lowerCamelCase__  =     name.replace("c_fc"				,     "fc1")
 if "c_proj" in name:
  lowerCamelCase__  =     name.replace("c_proj"				,     "fc2")
 if name.startswith("transformer.resblocks"):
  lowerCamelCase__  =     name.replace("transformer.resblocks"				,     "text_model.encoder.layers")
 if "attn.out_proj" in name and "message" not in name:
  lowerCamelCase__  =     name.replace("attn.out_proj"				,     "self_attn.out_proj")
 if "ln_final" in name:
  lowerCamelCase__  =     name.replace("ln_final"				,     "text_model.final_layer_norm")
 # visual encoder
 if name == "visual.class_embedding":
  lowerCamelCase__  =     name.replace("visual.class_embedding"				,     "vision_model.embeddings.class_embedding")
 if name == "visual.positional_embedding":
  lowerCamelCase__  =     name.replace("visual.positional_embedding"				,     "vision_model.embeddings.position_embedding.weight")
 if name.startswith("visual.transformer.resblocks"):
  lowerCamelCase__  =     name.replace("visual.transformer.resblocks"				,     "vision_model.encoder.layers")
 if "visual.conv1" in name:
  lowerCamelCase__  =     name.replace("visual.conv1"				,     "vision_model.embeddings.patch_embedding")
 if "visual.ln_pre" in name:
  lowerCamelCase__  =     name.replace("visual.ln_pre"				,     "vision_model.pre_layernorm")
 if "visual.ln_post" in name:
  lowerCamelCase__  =     name.replace("visual.ln_post"				,     "vision_model.post_layernorm")
 if "visual.proj" in name:
  lowerCamelCase__  =     name.replace("visual.proj"				,     "visual_projection.weight")
 if "text_projection" in name:
  lowerCamelCase__  =     name.replace("text_projection"				,     "text_projection.weight")
 # things on top
 if "prompts_visual_proj" in name:
  lowerCamelCase__  =     name.replace("prompts_visual_proj"				,     "prompts_visual_projection")
 if "prompts_visual_ln" in name:
  lowerCamelCase__  =     name.replace("prompts_visual_ln"				,     "prompts_visual_layernorm")
 # mit
 if name == "mit.positional_embedding":
  lowerCamelCase__  =     name.replace("positional"				,     "position")
 if name.startswith("mit.resblocks"):
  lowerCamelCase__  =     name.replace("mit.resblocks"				,     "mit.encoder.layers")
 # prompts generator
 if name.startswith("prompts_generator.norm"):
  lowerCamelCase__  =     name.replace("prompts_generator.norm"				,     "prompts_generator.layernorm")

 return name






def    lowerCamelCase_     (    lowercase__				,     lowercase__):
 for key in orig_state_dict.copy().keys():
  lowerCamelCase__  =     orig_state_dict.pop(lowercase__)

  if "attn.in_proj" in key:
   lowerCamelCase__  =     key.split(".")
   if key.startswith("visual"):
    lowerCamelCase__  =     key_split[3]
    lowerCamelCase__  =     config.vision_config.hidden_size
    if "message_attn" in key:
     if "weight" in key:
      lowerCamelCase__  =     val[
          :dim, :
      ]
      lowerCamelCase__  =     val[
          dim : dim * 2, :
      ]
      lowerCamelCase__  =     val[
          -dim:, :
      ]
     else:
      lowerCamelCase__  =     val[
          :dim
      ]
      lowerCamelCase__  =     val[
          dim : dim * 2
      ]
      lowerCamelCase__  =     val[
          -dim:
      ]
    else:
     if "weight" in key:
      lowerCamelCase__  =     val[
          :dim, :
      ]
      lowerCamelCase__  =     val[
          dim : dim * 2, :
      ]
      lowerCamelCase__  =     val[
          -dim:, :
      ]
     else:
      lowerCamelCase__  =     val[:dim]
      lowerCamelCase__  =     val[
          dim : dim * 2
      ]
      lowerCamelCase__  =     val[-dim:]
   elif key.startswith("mit"):
    lowerCamelCase__  =     key_split[2]
    lowerCamelCase__  =     config.vision_config.mit_hidden_size
    if "weight" in key:
     lowerCamelCase__  =     val[:dim, :]
     lowerCamelCase__  =     val[dim : dim * 2, :]
     lowerCamelCase__  =     val[-dim:, :]
    else:
     lowerCamelCase__  =     val[:dim]
     lowerCamelCase__  =     val[dim : dim * 2]
     lowerCamelCase__  =     val[-dim:]
   else:
    lowerCamelCase__  =     key_split[2]
    lowerCamelCase__  =     config.text_config.hidden_size
    if "weight" in key:
     lowerCamelCase__  =     val[:dim, :]
     lowerCamelCase__  =     val[
         dim : dim * 2, :
     ]
     lowerCamelCase__  =     val[-dim:, :]
    else:
     lowerCamelCase__  =     val[:dim]
     lowerCamelCase__  =     val[
         dim : dim * 2
     ]
     lowerCamelCase__  =     val[-dim:]
  else:
   lowerCamelCase__  =     rename_key(lowercase__)
   if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
    lowerCamelCase__  =     val.T
   lowerCamelCase__  =     val

 return orig_state_dict






def    lowerCamelCase_     (    lowercase__):
 if num_frames == 8:
  lowerCamelCase__  =     "eating_spaghetti_8_frames.npy"
 elif num_frames == 16:
  lowerCamelCase__  =     "eating_spaghetti.npy"
 elif num_frames == 32:
  lowerCamelCase__  =     "eating_spaghetti_32_frames.npy"
 lowerCamelCase__  =     hf_hub_download(
     repo_id="hf-internal-testing/spaghetti-video"				,     filename=lowercase__				,     repo_type="dataset"				,     )
 lowerCamelCase__  =     np.load(lowercase__)
 return list(lowercase__)






def    lowerCamelCase_     (    lowercase__				,     lowercase__=None				,     lowercase__=False):
 lowerCamelCase__  =     {
     # fully supervised kinetics-400 checkpoints
     "xclip-base-patch32": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth",
     "xclip-base-patch32-16-frames": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth"
     ),
     "xclip-base-patch16": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth",
     "xclip-base-patch16-16-frames": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth"
     ),
     "xclip-large-patch14": "https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb",
     "xclip-large-patch14-16-frames": "https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f",
     # fully supervised kinetics-600 checkpoints
     "xclip-base-patch16-kinetics-600": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth"
     ),
     "xclip-base-patch16-kinetics-600-16-frames": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth"
     ),
     "xclip-large-patch14-kinetics-600": "https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be",
     # few shot
     "xclip-base-patch16-hmdb-2-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth"
     ),
     "xclip-base-patch16-hmdb-4-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth"
     ),
     "xclip-base-patch16-hmdb-8-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth"
     ),
     "xclip-base-patch16-hmdb-16-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth"
     ),
     "xclip-base-patch16-ucf-2-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth"
     ),
     "xclip-base-patch16-ucf-4-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth"
     ),
     "xclip-base-patch16-ucf-8-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth"
     ),
     "xclip-base-patch16-ucf-16-shot": (
         "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth"
     ),
     # zero shot
     "xclip-base-patch16-zero-shot": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth",
 }

 lowerCamelCase__  =     model_to_url[model_name]
 lowerCamelCase__  =     8
 if "16-frames" in model_name:
  lowerCamelCase__  =     16
 elif "shot" in model_name:
  lowerCamelCase__  =     32

 lowerCamelCase__  =     get_xclip_config(lowercase__				,     lowercase__)
 lowerCamelCase__  =     XCLIPModel(lowercase__)
 model.eval()

 if "drive" in checkpoint_url:
  lowerCamelCase__  =     "pytorch_model.bin"
  gdown.cached_download(lowercase__				,     lowercase__				,     quiet=lowercase__)
  lowerCamelCase__  =     torch.load(lowercase__				,     map_location="cpu")["model"]
 else:
  lowerCamelCase__  =     torch.hub.load_state_dict_from_url(lowercase__)["model"]

 lowerCamelCase__  =     convert_state_dict(lowercase__				,     lowercase__)

 lowerCamelCase__  =     XCLIPModel(lowercase__)
 lowerCamelCase__  ,	lowerCamelCase__  =     model.load_state_dict(lowercase__				,     strict=lowercase__)
 assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
 model.eval()

 lowerCamelCase__  =     336 if model_name == "xclip-large-patch14-16-frames" else 224
 lowerCamelCase__  =     VideoMAEImageProcessor(size=lowercase__)
 lowerCamelCase__  =     CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
 lowerCamelCase__  =     CLIPTokenizerFast.from_pretrained("openai/clip-vit-base-patch32")
 lowerCamelCase__  =     XCLIPProcessor(image_processor=lowercase__				,     tokenizer=lowercase__)

 lowerCamelCase__  =     prepare_video(lowercase__)
 lowerCamelCase__  =     processor(
     text=["playing sports", "eating spaghetti", "go shopping"]				,     videos=lowercase__				,     return_tensors="pt"				,     padding=lowercase__)

 print("Shape of pixel values:"				,     inputs.pixel_values.shape)

 with torch.no_grad():
  lowerCamelCase__  =     model(**lowercase__)

 # Verify outputs
 lowerCamelCase__  =     outputs.logits_per_video
 lowerCamelCase__  =     logits_per_video.softmax(dim=1)
 print("Probs:"				,     lowercase__)
 # kinetics-400
 if model_name == "xclip-base-patch32":
  lowerCamelCase__  =     torch.tensor([[0.0_019, 0.9_951, 0.0_030]])
 elif model_name == "xclip-base-patch32-16-frames":
  lowerCamelCase__  =     torch.tensor([[7.0999E-04, 9.9883E-01, 4.5580E-04]])
 elif model_name == "xclip-base-patch16":
  lowerCamelCase__  =     torch.tensor([[0.0_083, 0.9_681, 0.0_236]])
 elif model_name == "xclip-base-patch16-16-frames":
  lowerCamelCase__  =     torch.tensor([[7.6937E-04, 9.9728E-01, 1.9473E-03]])
 elif model_name == "xclip-large-patch14":
  lowerCamelCase__  =     torch.tensor([[0.0_062, 0.9_864, 0.0_075]])
 elif model_name == "xclip-large-patch14-16-frames":
  lowerCamelCase__  =     torch.tensor([[3.3877E-04, 9.9937E-01, 2.8888E-04]])
 # kinetics-600
 elif model_name == "xclip-base-patch16-kinetics-600":
  lowerCamelCase__  =     torch.tensor([[0.0_555, 0.8_914, 0.0_531]])
 elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
  lowerCamelCase__  =     torch.tensor([[3.8554E-04, 9.9929E-01, 3.2754E-04]])
 elif model_name == "xclip-large-patch14-kinetics-600":
  lowerCamelCase__  =     torch.tensor([[0.0_036, 0.9_920, 0.0_045]])
 # few shot
 elif model_name == "xclip-base-patch16-hmdb-2-shot":
  lowerCamelCase__  =     torch.tensor([[7.1890E-06, 9.9994E-01, 5.6559E-05]])
 elif model_name == "xclip-base-patch16-hmdb-4-shot":
  lowerCamelCase__  =     torch.tensor([[1.0320E-05, 9.9993E-01, 6.2435E-05]])
 elif model_name == "xclip-base-patch16-hmdb-8-shot":
  lowerCamelCase__  =     torch.tensor([[4.1377E-06, 9.9990E-01, 9.8386E-05]])
 elif model_name == "xclip-base-patch16-hmdb-16-shot":
  lowerCamelCase__  =     torch.tensor([[4.1347E-05, 9.9962E-01, 3.3411E-04]])
 elif model_name == "xclip-base-patch16-ucf-2-shot":
  lowerCamelCase__  =     torch.tensor([[8.5857E-05, 9.9928E-01, 6.3291E-04]])
 elif model_name == "xclip-base-patch16-ucf-4-shot":
  lowerCamelCase__  =     torch.tensor([[8.5857E-05, 9.9928E-01, 6.3291E-04]])
 elif model_name == "xclip-base-patch16-ucf-8-shot":
  lowerCamelCase__  =     torch.tensor([[0.0_027, 0.9_904, 0.0_070]])
 elif model_name == "xclip-base-patch16-ucf-16-shot":
  lowerCamelCase__  =     torch.tensor([[9.8219E-04, 9.9593E-01, 3.0863E-03]])
 # zero shot
 elif model_name == "xclip-base-patch16-zero-shot":
  lowerCamelCase__  =     torch.tensor([[3.5082E-04, 9.9785E-01, 1.7966E-03]])
 else:
  raise ValueError(F'''Model name {model_name} not supported''')
 assert torch.allclose(lowercase__				,     lowercase__				,     atol=1E-3)
 print("Looks ok!")

 if pytorch_dump_folder_path is not None:
  print(F'''Saving model {model_name} to {pytorch_dump_folder_path}''')
  model.save_pretrained(lowercase__)

 if push_to_hub:
  print("Pushing model, processor and slow tokenizer files to the hub...")
  model.push_to_hub(lowercase__				,     organization="nielsr")
  processor.push_to_hub(lowercase__				,     organization="nielsr")
  slow_tokenizer.push_to_hub(lowercase__				,     organization="nielsr")


if __name__ == "__main__":
 __A							:							List[Any]  							=		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--model_name""",
     default="""xclip-base-patch32""",
     type=str,
     help="""Name of the model.""",
 )
 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."""
 )

 __A							:							str  							=		parser.parse_args()
 convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''





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

import numpy as np
import pyarrow as pa

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


if TYPE_CHECKING:
 import jax
 import jaxlib

__A							:							Any  							=		get_logger()

__A							:							Optional[dict]  							=		None
class 				lowercase       (							TensorFormatter[Mapping, "jax.Array", Mapping]      ):


    '''simple docstring'''

    def __init__(					self     :     str				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Dict=None				,     **__lowerCamelCase     :     str     )							->							List[Any]:

     '''simple docstring'''




     super().__init__(features=__lowerCamelCase     )
     import jax
     from jaxlib.xla_client import Device

     if isinstance(__lowerCamelCase				,     __lowerCamelCase     ):
      raise ValueError(
          f'''Expected {device} to be a `str` not {type(__lowerCamelCase     )}, as `jaxlib.xla_extension.Device` '''
          "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
          "the device with `str()` to get its string identifier that will be internally mapped "
          "to the actual `jaxlib.xla_extension.Device`."     )
     lowerCamelCase__  =     device if isinstance(__lowerCamelCase				,     __lowerCamelCase     ) else str(jax.devices()[0]     )
     # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
     # with `pickle` nor with `dill`, so we need to use a global variable instead
     global DEVICE_MAPPING
     if DEVICE_MAPPING is None:
      lowerCamelCase__  =     self._map_devices_to_str()
     if self.device not in list(DEVICE_MAPPING.keys()     ):
      logger.warning(
          f'''Device with string identifier {self.device} not listed among the available '''
          f'''devices: {list(DEVICE_MAPPING.keys()     )}, so falling back to the default '''
          f'''device: {str(jax.devices()[0]     )}.'''     )
      lowerCamelCase__  =     str(jax.devices()[0]     )
     lowerCamelCase__  =     jnp_array_kwargs

    @staticmethod
    def    a__				(					)							->							Dict[str, "jaxlib.xla_extension.Device"]:

     '''simple docstring'''




     import jax

     return {str(__lowerCamelCase     ): device for device in jax.devices()}

    def    a__				(					self     :     Optional[Any]				,     __lowerCamelCase     :     Optional[int]     )							->							List[str]:

     '''simple docstring'''




     import jax
     import jax.numpy as jnp

     if isinstance(__lowerCamelCase				,     __lowerCamelCase     ) and column:
      if all(
          isinstance(__lowerCamelCase				,     jax.Array     ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column     ):
       return jnp.stack(__lowerCamelCase				,     axis=0     )
     return column

    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     List[Any]     )							->							Optional[int]:

     '''simple docstring'''




     import jax
     import jax.numpy as jnp

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

     lowerCamelCase__  =     {}

     if isinstance(__lowerCamelCase				,     (np.number, np.ndarray)     ) and np.issubdtype(value.dtype				,     np.integer     ):
      # the default int precision depends on the jax config
      # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
      if jax.config.jax_enable_xaa:
       lowerCamelCase__  =     {"dtype": jnp.intaa}
      else:
       lowerCamelCase__  =     {"dtype": jnp.intaa}
     elif isinstance(__lowerCamelCase				,     (np.number, np.ndarray)     ) and np.issubdtype(value.dtype				,     np.floating     ):
      lowerCamelCase__  =     {"dtype": jnp.floataa}
     elif config.PIL_AVAILABLE and "PIL" in sys.modules:
      import PIL.Image

      if isinstance(__lowerCamelCase				,     PIL.Image.Image     ):
       lowerCamelCase__  =     np.asarray(__lowerCamelCase     )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
     global DEVICE_MAPPING
     if DEVICE_MAPPING is None:
      lowerCamelCase__  =     self._map_devices_to_str()

     with jax.default_device(DEVICE_MAPPING[self.device]     ):
      # calling jnp.array on a np.ndarray does copy the data
      # see https://github.com/google/jax/issues/4486
      return jnp.array(__lowerCamelCase				,     **{**default_dtype, **self.jnp_array_kwargs}     )

    def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     Optional[Any]     )							->							Dict:

     '''simple docstring'''




     import jax

     # support for torch, tf, jax etc.
     if config.TORCH_AVAILABLE and "torch" in sys.modules:
      import torch

      if isinstance(__lowerCamelCase				,     torch.Tensor     ):
       return self._tensorize(data_struct.detach().cpu().numpy()[()]     )
     if hasattr(__lowerCamelCase				,     "__array__"     ) and not isinstance(__lowerCamelCase				,     jax.Array     ):
      lowerCamelCase__  =     data_struct.__array__()
     # support for nested types like struct of list of struct
     if isinstance(__lowerCamelCase				,     np.ndarray     ):
      if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
       return self._consolidate([self.recursive_tensorize(__lowerCamelCase     ) for substruct in data_struct]     )
     elif isinstance(__lowerCamelCase				,     (list, tuple)     ):
      return self._consolidate([self.recursive_tensorize(__lowerCamelCase     ) for substruct in data_struct]     )
     return self._tensorize(__lowerCamelCase     )

    def    a__				(					self     :     Any				,     __lowerCamelCase     :     dict     )							->							int:

     '''simple docstring'''




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

    def    a__				(					self     :     Union[str, Any]				,     __lowerCamelCase     :     pa.Table     )							->							Mapping:

     '''simple docstring'''




     lowerCamelCase__  =     self.numpy_arrow_extractor().extract_row(__lowerCamelCase     )
     lowerCamelCase__  =     self.python_features_decoder.decode_row(__lowerCamelCase     )
     return self.recursive_tensorize(__lowerCamelCase     )

    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     pa.Table     )							->							"jax.Array":

     '''simple docstring'''




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




    def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     pa.Table     )							->							Mapping:

     '''simple docstring'''




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

