





# 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

_snake_case :    Optional[int]           =			get_logger()

_snake_case :    str           =			None





class  lowerCAmelCase							(       TensorFormatter[Mapping, 'jax.Array', Mapping] ):







  def __init__(      self  :Dict   ,							_lowercase  :int=None   ,							_lowercase  :Tuple=None   ,							**_lowercase  :str ):






     '''simple docstring'''





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

     if isinstance(_lowercase   ,							_lowercase ):
        raise ValueError(
            f'''Expected {device} to be a `str` not {type(_lowercase )}, 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`." )
     lowercase__			       =  device if isinstance(_lowercase   ,							_lowercase ) 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:
        lowercase__			       =  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] )}.''' )
        lowercase__			       =  str(jax.devices()[0] )
     lowercase__			       =  jnp_array_kwargs







  @staticmethod
  def   UpperCAmelCase					(      ):






     '''simple docstring'''





     import jax

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







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Tuple ):






     '''simple docstring'''





     import jax
     import jax.numpy as jnp

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







  def   UpperCAmelCase					(      self  :str   ,							_lowercase  :int ):






     '''simple docstring'''





     import jax
     import jax.numpy as jnp

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

     lowercase__			       =  {}

     if isinstance(_lowercase   ,							(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:
           lowercase__			       =  {"dtype": jnp.intaa}
        else:
           lowercase__			       =  {"dtype": jnp.intaa}
     elif isinstance(_lowercase   ,							(np.number, np.ndarray) ) and np.issubdtype(value.dtype   ,							np.floating ):
        lowercase__			       =  {"dtype": jnp.floataa}
     elif config.PIL_AVAILABLE and "PIL" in sys.modules:
        import PIL.Image

        if isinstance(_lowercase   ,							PIL.Image.Image ):
           lowercase__			       =  np.asarray(_lowercase )

        # 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:
        lowercase__			       =  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(_lowercase   ,							**{**default_dtype, **self.jnp_array_kwargs} )







  def   UpperCAmelCase					(      self  :Any   ,							_lowercase  :Any ):






     '''simple docstring'''





     import jax

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

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







  def   UpperCAmelCase					(      self  :List[Any]   ,							_lowercase  :dict ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :pa.Table ):






     '''simple docstring'''





     lowercase__			       =  self.numpy_arrow_extractor().extract_row(_lowercase )
     lowercase__			       =  self.python_features_decoder.decode_row(_lowercase )
     return self.recursive_tensorize(_lowercase )







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :pa.Table ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Dict   ,							_lowercase  :pa.Table ):






     '''simple docstring'''





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













from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml





class  lowerCAmelCase							(       yaml.SafeLoader ):







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :Any ):






     '''simple docstring'''





     lowercase__			       =  [self.constructed_objects[key_node] for key_node, _ in node.value]
     lowercase__			       =  [tuple(_lowercase ) if isinstance(_lowercase   ,							_lowercase ) else key for key in keys]
     lowercase__			       =  Counter(_lowercase )
     lowercase__			       =  [key for key in counter if counter[key] > 1]
     if duplicate_keys:
        raise TypeError(f'''Got duplicate yaml keys: {duplicate_keys}''' )







  def   UpperCAmelCase					(      self  :Any   ,							_lowercase  :str   ,							_lowercase  :Dict=False ):






     '''simple docstring'''





     lowercase__			       =  super().construct_mapping(_lowercase   ,							deep=_lowercase )
     self._check_no_duplicates_on_constructed_node(_lowercase )
     return mapping


def _A       (   __magic_name__	):
   lowercase__			       =  list(readme_content.splitlines()	)
   if full_content and full_content[0] == "---" and "---" in full_content[1:]:
      lowercase__			       =  full_content[1:].index("---"	) + 1
      lowercase__			       =  "\n".join(full_content[1:sep_idx]	)
      return yamlblock, "\n".join(full_content[sep_idx + 1 :]	)

   return None, "\n".join(__magic_name__	)





class  lowerCAmelCase							(       lowercase_ ):
  # class attributes
  __lowerCamelCase		=      {'train_eval_index'}  # train-eval-index in the YAML metadata







  @classmethod
  def   UpperCAmelCase					(      cls  :Dict   ,							_lowercase  :Path ):






     '''simple docstring'''





     with open(_lowercase   ,							encoding="utf-8" ) as readme_file:
        lowercase__       ,							lowercase__			       =  _split_yaml_from_readme(readme_file.read() )
     if yaml_string is not None:
        return cls.from_yaml_string(_lowercase )
     else:
        return cls()







  def   UpperCAmelCase					(      self  :Any   ,							_lowercase  :Path ):






     '''simple docstring'''





     if path.exists():
        with open(_lowercase   ,							encoding="utf-8" ) as readme_file:
           lowercase__			       =  readme_file.read()
     else:
        lowercase__			       =  None
     lowercase__			       =  self._to_readme(_lowercase )
     with open(_lowercase   ,							"w"   ,							encoding="utf-8" ) as readme_file:
        readme_file.write(_lowercase )







  def   UpperCAmelCase					(      self  :List[str]   ,							_lowercase  :Optional[str] = None ):






     '''simple docstring'''





     if readme_content is not None:
        lowercase__       ,							lowercase__			       =  _split_yaml_from_readme(_lowercase )
        lowercase__			       =  "---\n" + self.to_yaml_string() + "---\n" + content
     else:
        lowercase__			       =  "---\n" + self.to_yaml_string() + "---\n"
     return full_content







  @classmethod
  def   UpperCAmelCase					(      cls  :Any   ,							_lowercase  :str ):






     '''simple docstring'''





     lowercase__			       =  yaml.load(_lowercase   ,							Loader=_NoDuplicateSafeLoader ) or {}

     # Convert the YAML keys to DatasetMetadata fields
     lowercase__			       =  {
         (key.replace("-"   ,							"_" ) if key.replace("-"   ,							"_" ) in cls._FIELDS_WITH_DASHES else key): value
         for key, value in metadata_dict.items()
     }
     return cls(**_lowercase )







  def   UpperCAmelCase					(      self  :Union[str, Any] ):






     '''simple docstring'''





     return yaml.safe_dump(
         {
             (key.replace("_"   ,							"-" ) if key in self._FIELDS_WITH_DASHES else key): value
             for key, value in self.items()
         }   ,							sort_keys=_lowercase   ,							allow_unicode=_lowercase   ,							encoding="utf-8"   ,							).decode("utf-8" )
_snake_case           =			{
    """image-classification""": [],
    """translation""": [],
    """image-segmentation""": [],
    """fill-mask""": [],
    """automatic-speech-recognition""": [],
    """token-classification""": [],
    """sentence-similarity""": [],
    """audio-classification""": [],
    """question-answering""": [],
    """summarization""": [],
    """zero-shot-classification""": [],
    """table-to-text""": [],
    """feature-extraction""": [],
    """other""": [],
    """multiple-choice""": [],
    """text-classification""": [],
    """text-to-image""": [],
    """text2text-generation""": [],
    """zero-shot-image-classification""": [],
    """tabular-classification""": [],
    """tabular-regression""": [],
    """image-to-image""": [],
    """tabular-to-text""": [],
    """unconditional-image-generation""": [],
    """text-retrieval""": [],
    """text-to-speech""": [],
    """object-detection""": [],
    """audio-to-audio""": [],
    """text-generation""": [],
    """conversational""": [],
    """table-question-answering""": [],
    """visual-question-answering""": [],
    """image-to-text""": [],
    """reinforcement-learning""": [],
    """voice-activity-detection""": [],
    """time-series-forecasting""": [],
    """document-question-answering""": [],
}


if __name__ == "__main__":
  from argparse import ArgumentParser

  _snake_case           =			ArgumentParser(usage="""Validate the yaml metadata block of a README.md file.""")
  ap.add_argument("""readme_filepath""")
  _snake_case           =			ap.parse_args()

  _snake_case           =			Path(args.readme_filepath)
  _snake_case           =			DatasetMetadata.from_readme(readme_filepath)
  print(dataset_metadata)
  dataset_metadata.to_readme(readme_filepath)






