



import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


__lowerCamelCase         =			datasets.utils.logging.get_logger(__name__)


@dataclass
class   snake_case_ (datasets.BuilderConfig  ):




       """simple docstring"""






       _lowerCamelCase			    =				1_0000
       _lowerCamelCase			    =				None
       _lowerCamelCase			    =				None


class   snake_case_ (datasets.ArrowBasedBuilder  ):




       """simple docstring"""






       _lowerCamelCase			    =				ParquetConfig



       def  A_   (     self):

             """simple docstring"""





             return datasets.DatasetInfo(features=self.config.features)



       def  A_   (     self     ,lowercase):

             """simple docstring"""





             if not self.config.data_files:
                   raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}""")
             UpperCAmelCase_  :	List[Any]				      =		dl_manager.download_and_extract(self.config.data_files)
             if isinstance(lowercase     ,(str, list, tuple)):
                   UpperCAmelCase_  :	Any				      =		data_files
                   if isinstance(lowercase     ,lowercase):
                         UpperCAmelCase_  :	List[str]				      =		[files]
                   # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
                   UpperCAmelCase_  :	Any				      =		[dl_manager.iter_files(lowercase) for file in files]
                   return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,gen_kwargs={"files": files})]
             UpperCAmelCase_  :	List[Any]				      =		[]
             for split_name, files in data_files.items():
                   if isinstance(lowercase     ,lowercase):
                         UpperCAmelCase_  :	Tuple				      =		[files]
                   # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
                   UpperCAmelCase_  :	Union[str, Any]				      =		[dl_manager.iter_files(lowercase) for file in files]
                   # Infer features is they are stoed in the arrow schema
                   if self.info.features is None:
                         for file in itertools.chain.from_iterable(lowercase):
                               with open(lowercase     ,"rb") as f:
                                     UpperCAmelCase_  :	List[Any]				      =		datasets.Features.from_arrow_schema(pq.read_schema(lowercase))
                               break
                   splits.append(datasets.SplitGenerator(name=lowercase     ,gen_kwargs={"files": files}))
             return splits



       def  A_   (     self     ,lowercase):

             """simple docstring"""





             if self.info.features is not None:
                   # more expensive cast to support nested features with keys in a different order
                   # allows str <-> int/float or str to Audio for example
                   UpperCAmelCase_  :	int				      =		table_cast(lowercase     ,self.info.features.arrow_schema)
             return pa_table



       def  A_   (     self     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		self.info.features.arrow_schema if self.info.features is not None else None
             if self.info.features is not None and self.config.columns is not None:
                   if sorted(field.name for field in schema) != sorted(self.config.columns):
                         raise ValueError(
                             F"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'""")
             for file_idx, file in enumerate(itertools.chain.from_iterable(lowercase)):
                   with open(lowercase     ,"rb") as f:
                         UpperCAmelCase_  :	List[Any]				      =		pq.ParquetFile(lowercase)
                         try:
                               for batch_idx, record_batch in enumerate(
                                   parquet_file.iter_batches(batch_size=self.config.batch_size     ,columns=self.config.columns)):
                                     UpperCAmelCase_  :	int				      =		pa.Table.from_batches([record_batch])
                                     # Uncomment for debugging (will print the Arrow table size and elements)
                                     # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                     # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                     yield F"""{file_idx}_{batch_idx}""", self._cast_table(lowercase)
                         except ValueError as e:
                               logger.error(F"""Failed to read file '{file}' with error {type(lowercase)}: {e}""")
                               raise





from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''google/umt5-small''': '''https://huggingface.co/google/umt5-small/resolve/main/config.json''',
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""umt5"""
       _lowerCamelCase			    =				["""past_key_values"""]



       def __init__(     self     ,lowercase=250112     ,lowercase=512     ,lowercase=64     ,lowercase=1024     ,lowercase=8     ,lowercase=None     ,lowercase=6     ,lowercase=32     ,lowercase=128     ,lowercase=0.1     ,lowercase=1E-6     ,lowercase=1.0     ,lowercase="gated-gelu"     ,lowercase=True     ,lowercase=True     ,lowercase="T5Tokenizer"     ,lowercase=True     ,lowercase=0     ,lowercase=1     ,lowercase=0     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(
                 is_encoder_decoder=lowercase     ,tokenizer_class=lowercase     ,tie_word_embeddings=lowercase     ,pad_token_id=lowercase     ,eos_token_id=lowercase     ,decoder_start_token_id=lowercase     ,**lowercase     ,)
             UpperCAmelCase_  :	Optional[int]				      =		vocab_size
             UpperCAmelCase_  :	Any				      =		d_model
             UpperCAmelCase_  :	Any				      =		d_kv
             UpperCAmelCase_  :	int				      =		d_ff
             UpperCAmelCase_  :	Tuple				      =		num_layers
             UpperCAmelCase_  :	int				      =		(
                 num_decoder_layers if num_decoder_layers is not None else self.num_layers
             )  # default = symmetry
             UpperCAmelCase_  :	Optional[int]				      =		num_heads
             UpperCAmelCase_  :	str				      =		relative_attention_num_buckets
             UpperCAmelCase_  :	Any				      =		relative_attention_max_distance
             UpperCAmelCase_  :	Optional[Any]				      =		dropout_rate
             UpperCAmelCase_  :	Union[str, Any]				      =		layer_norm_epsilon
             UpperCAmelCase_  :	Optional[Any]				      =		initializer_factor
             UpperCAmelCase_  :	int				      =		feed_forward_proj
             UpperCAmelCase_  :	str				      =		use_cache

             UpperCAmelCase_  :	List[str]				      =		self.feed_forward_proj.split("-")
             UpperCAmelCase_  :	Any				      =		act_info[-1]
             UpperCAmelCase_  :	Optional[int]				      =		act_info[0] == "gated"

             if len(lowercase) > 1 and act_info[0] != "gated" or len(lowercase) > 2:
                   raise ValueError(
                       F"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
                       "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
                       "'gated-gelu' or 'relu'")

             if feed_forward_proj == "gated-gelu":
                   UpperCAmelCase_  :	Tuple				      =		"gelu_new"



       @property
       def  A_   (     self):

             """simple docstring"""





             return self.d_model



       @property
       def  A_   (     self):

             """simple docstring"""





             return self.num_heads



       @property
       def  A_   (     self):

             """simple docstring"""





             return self.num_layers


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       @property
       # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		{
                 "input_ids": {0: "batch", 1: "encoder_sequence"},
                 "attention_mask": {0: "batch", 1: "encoder_sequence"},
             }
             if self.use_past:
                   UpperCAmelCase_  :	Union[str, Any]				      =		"past_encoder_sequence + sequence"
                   UpperCAmelCase_  :	Optional[int]				      =		{0: "batch"}
                   UpperCAmelCase_  :	Union[str, Any]				      =		{0: "batch", 1: "past_decoder_sequence + sequence"}
             else:
                   UpperCAmelCase_  :	Optional[int]				      =		{0: "batch", 1: "decoder_sequence"}
                   UpperCAmelCase_  :	Dict				      =		{0: "batch", 1: "decoder_sequence"}

             if self.use_past:
                   self.fill_with_past_key_values_(lowercase     ,direction="inputs")

             return common_inputs



       @property
       # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
       def  A_   (     self):

             """simple docstring"""





             return 13



       @property
       def  A_   (     self):

             """simple docstring"""





             return 5E-4
