





"""simple docstring"""


# 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 tensorflow as tf



class    __SCREAMING_SNAKE_CASE			(  TensorFormatter[Mapping, "tf.Tensor", Mapping]      ):







  '''simple docstring'''


  def __init__(						self      :				Union[str, Any]    ,      __a      :				int=None    ,      **__a      :				Tuple		)						-> List[Any]:
        super().__init__(features=__a		)
        _UpperCamelCase					:		Any  =				tf_tensor_kwargs
        import tensorflow as tf  # noqa: import tf at initialization

  def       __SCREAMING_SNAKE_CASE (						self      :				List[str]    ,      __a      :				List[str]		)						-> Any:
        import tensorflow as tf

        if isinstance(__a    ,      __a		) and column:
              if all(
                  isinstance(__a    ,      tf.Tensor		) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column		):
                    return tf.stack(__a		)
              elif all(
                  isinstance(__a    ,      (tf.Tensor, tf.RaggedTensor)		) and x.ndim == 1 and x.dtype == column[0].dtype
                  for x in column		):
                    # only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
                    return tf.ragged.stack(__a		)

        return column

  def       __SCREAMING_SNAKE_CASE (						self      :				Union[str, Any]    ,      __a      :				Any		)						-> Dict:
        import tensorflow as tf

        if value is None:
              return value

        _UpperCamelCase					:		str  =				{}

        if isinstance(__a    ,      (np.number, np.ndarray)		) and np.issubdtype(value.dtype    ,      np.integer		):
              _UpperCamelCase					:		int  =				{"dtype": tf.intaa}
        elif isinstance(__a    ,      (np.number, np.ndarray)		) and np.issubdtype(value.dtype    ,      np.floating		):
              _UpperCamelCase					:		Any  =				{"dtype": tf.floataa}
        elif config.PIL_AVAILABLE and "PIL" in sys.modules:
              import PIL.Image

              if isinstance(__a    ,      PIL.Image.Image		):
                    _UpperCamelCase					:		Tuple  =				np.asarray(__a		)

        return tf.convert_to_tensor(__a    ,      **{**default_dtype, **self.tf_tensor_kwargs}		)

  def       __SCREAMING_SNAKE_CASE (						self      :				Any    ,      __a      :				Tuple		)						-> str:
        import tensorflow as tf

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

              if isinstance(__a    ,      torch.Tensor		):
                    return self._tensorize(data_struct.detach().cpu().numpy()[()]		)
        if hasattr(__a    ,      "__array__"		) and not isinstance(__a    ,      tf.Tensor		):
              _UpperCamelCase					:		str  =				data_struct.__array__()
        # support for nested types like struct of list of struct
        if isinstance(__a    ,      np.ndarray		):
              if data_struct.dtype == object:  # tf 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       __SCREAMING_SNAKE_CASE (						self      :				Any    ,      __a      :				dict		)						-> Any:
        return map_nested(self._recursive_tensorize    ,      __a    ,      map_list=__a		)

  def       __SCREAMING_SNAKE_CASE (						self      :				int    ,      __a      :				pa.Table		)						-> Mapping:
        _UpperCamelCase					:		Optional[int]  =				self.numpy_arrow_extractor().extract_row(__a		)
        _UpperCamelCase					:		Optional[int]  =				self.python_features_decoder.decode_row(__a		)
        return self.recursive_tensorize(__a		)

  def       __SCREAMING_SNAKE_CASE (						self      :				Optional[int]    ,      __a      :				pa.Table		)						-> "tf.Tensor":
        _UpperCamelCase					:		Optional[int]  =				self.numpy_arrow_extractor().extract_column(__a		)
        _UpperCamelCase					:		str  =				self.python_features_decoder.decode_column(__a    ,      pa_table.column_names[0]		)
        _UpperCamelCase					:		Optional[Any]  =				self.recursive_tensorize(__a		)
        _UpperCamelCase					:		List[str]  =				self._consolidate(__a		)
        return column


  def       __SCREAMING_SNAKE_CASE (						self      :				List[Any]    ,      __a      :				pa.Table		)						-> Mapping:
        _UpperCamelCase					:		List[str]  =				self.numpy_arrow_extractor().extract_batch(__a		)
        _UpperCamelCase					:		Any  =				self.python_features_decoder.decode_batch(__a		)
        _UpperCamelCase					:		int  =				self.recursive_tensorize(__a		)
        for column_name in batch:
              _UpperCamelCase					:		Dict  =				self._consolidate(batch[column_name]		)
        return batch







"""simple docstring"""


import glob
import os

from get_test_info import get_tester_classes


if __name__ == "__main__":
  lowerCamelCase__      =      []

  lowerCamelCase__      =      os.path.join("tests", "models", "**", "test_modeling_*.py")
  lowerCamelCase__      =      glob.glob(pattern)
  # TODO: deal with TF/Flax too
  lowerCamelCase__      =      [
      x for x in test_files if not (x.startswith("test_modeling_tf_") or x.startswith("test_modeling_flax_"))
  ]

  for test_file in test_files:
    lowerCamelCase__      =      get_tester_classes(test_file)
    for tester_class in tester_classes:
      # A few tester classes don't have `parent` parameter in `__init__`.
      # TODO: deal this better
      try:
        lowerCamelCase__      =      tester_class(parent=None)
      except Exception:
        continue
      if hasattr(tester, "get_config"):
        lowerCamelCase__      =      tester.get_config()
        for k, v in config.to_dict().items():
          if isinstance(v, int):
            lowerCamelCase__      =      None
            if k in ["vocab_size"]:
              lowerCamelCase__      =      100
            elif k in ["max_position_embeddings"]:
              lowerCamelCase__      =      128
            elif k in ["hidden_size", "d_model"]:
              lowerCamelCase__      =      40
            elif k == ["num_layers", "num_hidden_layers", "num_encoder_layers", "num_decoder_layers"]:
              lowerCamelCase__      =      5
            if target is not None and v > target:
              failures.append(
                  f"""{tester_class.__name__} will produce a `config` of type `{config.__class__.__name__}`"""
                  f""" with config[\"{k}\"] = {v} which is too large for testing! Set its value to be smaller"""
                  f""" than {target}."""
              )

  if len(failures) > 0:
    raise Exception(f"""There were {len(failures)} failures:\n""" + "\n".join(failures))
