"""simple docstring"""






import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput


class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Any					=      dataset
       lowercase__:   Dict					=      process
       lowercase__:   List[Any]					=      params






    def __len__(   self							):
       return len(self.dataset							)






    def __getitem__(   self    ,			_UpperCAmelCase							):
       lowercase__:   List[str]					=      self.dataset[i]
       lowercase__:   Optional[Any]					=      self.process(_UpperCAmelCase    ,			**self.params							)
       return processed



class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase=None							):
       lowercase__:   Union[str, Any]					=      loader
       lowercase__:   Dict					=      infer
       lowercase__:   str					=      params
       if loader_batch_size == 1:
          # Let's spare some time by deactivating altogether
          lowercase__:   Any					=      None
       lowercase__:   Any					=      loader_batch_size

       # Internal bookkeeping
       lowercase__:   Any					=      None
       lowercase__:   List[str]					=      None






    def __len__(   self							):
       return len(self.loader							)






    def __iter__(   self							):
       lowercase__:   Tuple					=      iter(self.loader							)
       return self






    def _snake_case			(   self							):
       if isinstance(self._loader_batch_data    ,			torch.Tensor							):
          # Batch data is simple tensor, just fetch the slice
          lowercase__:   int					=      self._loader_batch_data[self._loader_batch_index]
       else:
          # Batch data is assumed to be BaseModelOutput (or dict)
          lowercase__:   Dict					=      {}
          for k, element in self._loader_batch_data.items():
             if isinstance(_UpperCAmelCase    ,			_UpperCAmelCase							):
                # Convert ModelOutput to tuple first
                lowercase__:   Union[str, Any]					=      element.to_tuple()
                if isinstance(element[0]    ,			torch.Tensor							):
                   lowercase__:   Any					=      tuple(el[self._loader_batch_index].unsqueeze(0							) for el in element							)
                elif isinstance(element[0]    ,			np.ndarray							):
                   lowercase__:   Optional[int]					=      tuple(np.expand_dims(el[self._loader_batch_index]    ,			0							) for el in element							)
                continue
             if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(_UpperCAmelCase    ,			_UpperCAmelCase							):
                # Those are stored as lists of tensors so need specific unbatching.
                if isinstance(element[0]    ,			torch.Tensor							):
                   lowercase__:   List[str]					=      tuple(el[self._loader_batch_index].unsqueeze(0							) for el in element							)
                elif isinstance(element[0]    ,			np.ndarray							):
                   lowercase__:   str					=      tuple(np.expand_dims(el[self._loader_batch_index]    ,			0							) for el in element							)
                continue
             if element is None:
                # This can happen for optional data that get passed around
                lowercase__:   Optional[int]					=      None
             elif isinstance(element[self._loader_batch_index]    ,			torch.Tensor							):
                # Take correct batch data, but make it looked like batch_size=1
                # For compatibility with other methods within transformers

                lowercase__:   Any					=      element[self._loader_batch_index].unsqueeze(0							)
             elif isinstance(element[self._loader_batch_index]    ,			np.ndarray							):
                # Take correct batch data, but make it looked like batch_size=1
                # For compatibility with other methods within transformers
                lowercase__:   List[str]					=      np.expand_dims(element[self._loader_batch_index]    ,			0							)
             else:
                # This is typically a list, so no need to `unsqueeze`.
                lowercase__:   List[Any]					=      element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
          lowercase__:   Any					=      self._loader_batch_data.__class__(_UpperCAmelCase							)
       self._loader_batch_index += 1
       return result






    def _snake_case			(   self							):
       if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
          # We are currently unrolling a batch so we just need to return
          # the current item within a batch
          return self.loader_batch_item()

       # We're out of items within a batch
       lowercase__:   Dict					=      next(self.iterator							)
       lowercase__:   Any					=      self.infer(_UpperCAmelCase    ,			**self.params							)
       # We now have a batch of "inferred things".
       if self.loader_batch_size is not None:
          # Try to infer the size of the batch
          if isinstance(_UpperCAmelCase    ,			torch.Tensor							):
             lowercase__:   Union[str, Any]					=      processed
          else:
             lowercase__:   Dict					=      list(processed.keys()							)[0]
             lowercase__:   Union[str, Any]					=      processed[key]
          if isinstance(_UpperCAmelCase    ,			_UpperCAmelCase							):
             lowercase__:   Any					=      len(_UpperCAmelCase							)
          else:
             lowercase__:   str					=      first_tensor.shape[0]
          if 0 < observed_batch_size < self.loader_batch_size:
             # could be last batch so we can't unroll as many
             # elements.
             lowercase__:   Union[str, Any]					=      observed_batch_size
          # Setting internal index to unwrap the batch
          lowercase__:   Tuple					=      processed
          lowercase__:   Optional[int]					=      0
          return self.loader_batch_item()
       else:
          # We're not unrolling batches
          return processed



class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase=None							):
       super().__init__(_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							)






    def __iter__(   self							):
       lowercase__:   Any					=      iter(self.loader							)
       lowercase__:   Optional[Any]					=      None
       return self






    def _snake_case			(   self							):
       if self.subiterator is None:
          lowercase__:   Optional[int]					=      self.infer(next(self.iterator							)    ,			**self.params							)
       try:
          # Try to return next item
          lowercase__:   str					=      next(self.subiterator							)
       except StopIteration:
          # When a preprocess iterator ends, we can start lookig at the next item
          # ChunkIterator will keep feeding until ALL elements of iterator
          # all have created their subiterator and have been iterating against.
          #
          # Another way to look at it, is we're basically flattening lists of lists
          # into a single list, but with generators
          lowercase__:   Any					=      self.infer(next(self.iterator							)    ,			**self.params							)
          lowercase__:   int					=      next(self.subiterator							)
       return processed



class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __iter__(   self							):
       lowercase__:   Optional[int]					=      iter(self.loader							)
       return self






    def _snake_case			(   self							):
       # Extremely similar to PipelineIterator in its unpacking mechanism
       # BUT, we have an extra required item which is the presence of `is_last`
       # That is because everything is flattened by `PipelineChunkIterator` we
       # need to keep track of how to regroup here in the original `process`
       # boundaries so that `process` and `postprocess` see the same data.

       # This iterator accumulates items (possibly while unbatching) until it
       # its a `is_last` and then just passes it on to the caller.
       lowercase__:   int					=      False
       lowercase__:   Any					=      []
       if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
          while self._loader_batch_index < self.loader_batch_size:
             lowercase__:   Union[str, Any]					=      self.loader_batch_item()
             lowercase__:   Tuple					=      item.pop('''is_last'''							)
             accumulator.append(_UpperCAmelCase							)
             if is_last:
                return accumulator

       while not is_last:
          lowercase__:   str					=      self.infer(next(self.iterator							)    ,			**self.params							)
          if self.loader_batch_size is not None:
             if isinstance(_UpperCAmelCase    ,			torch.Tensor							):
                lowercase__:   str					=      processed
             else:
                lowercase__:   Union[str, Any]					=      list(processed.keys()							)[0]
                lowercase__:   Optional[Any]					=      processed[key]
             if isinstance(_UpperCAmelCase    ,			_UpperCAmelCase							):
                lowercase__:   Tuple					=      len(_UpperCAmelCase							)
             else:
                lowercase__:   Optional[Any]					=      first_tensor.shape[0]
             if 0 < observed_batch_size < self.loader_batch_size:
                # could be last batch so we can't unroll as many
                # elements.
                lowercase__:   Any					=      observed_batch_size
             lowercase__:   int					=      processed
             lowercase__:   Dict					=      0
             while self._loader_batch_index < self.loader_batch_size:
                lowercase__:   Union[str, Any]					=      self.loader_batch_item()
                lowercase__:   Union[str, Any]					=      item.pop('''is_last'''							)
                accumulator.append(_UpperCAmelCase							)
                if is_last:
                   return accumulator
          else:
             lowercase__:   Union[str, Any]					=      processed
             lowercase__:   List[str]					=      item.pop('''is_last'''							)
             accumulator.append(_UpperCAmelCase							)
       return accumulator



class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Union[str, Any]					=      dataset
       lowercase__:   str					=      key






    def __len__(   self							):
       return len(self.dataset							)






    def __getitem__(   self    ,			_UpperCAmelCase							):
       return self.dataset[i][self.key]



class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Optional[int]					=      dataset
       lowercase__:   List[str]					=      keya
       lowercase__:   List[str]					=      keya






    def __len__(   self							):
       return len(self.dataset							)






    def __getitem__(   self    ,			_UpperCAmelCase							):
       return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}

"""simple docstring"""






import argparse

import torch

from transformers import FunnelBaseModel, FunnelConfig, FunnelModel, load_tf_weights_in_funnel
from transformers.utils import logging


logging.set_verbosity_info()



def SCREAMING_SNAKE_CASE__			(   __UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase )  ->  List[Any]:
   # Initialise PyTorch model
   lowercase__:   Optional[Any]					=      FunnelConfig.from_json_file(__UpperCAmelCase )
   print(F"""Building PyTorch model from configuration: {config}""" )
   lowercase__:   List[Any]					=      FunnelBaseModel(__UpperCAmelCase ) if base_model else FunnelModel(__UpperCAmelCase )

   # Load weights from tf checkpoint
   load_tf_weights_in_funnel(__UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase )

   # Save pytorch-model
   print(F"""Save PyTorch model to {pytorch_dump_path}""" )
   torch.save(model.state_dict()       ,    __UpperCAmelCase )


if __name__ == "__main__":
       __A			   = argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
       )
       parser.add_argument(
           "--config_file",
           default=None,
           type=str,
           required=True,
           help="The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.",
       )
       parser.add_argument(
           "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
       )
       parser.add_argument(
           "--base_model", action="store_true", help="Whether you want just the base model (no decoder) or not."
       )
       __A			   = parser.parse_args()
       convert_tf_checkpoint_to_pytorch(
           args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.base_model
       )
