'''simple docstring'''


from collections import OrderedDict
from typing import Mapping

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


snake_case__	:     List[str]      =    logging.get_logger(__name__)

snake_case__	:     Any      =    {
    '''mnaylor/mega-base-wikitext''': '''https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json''',
}






class 	__SCREAMING_SNAKE_CASE		(			lowerCamelCase_			):







  '''simple docstring'''






  lowerCamelCase_						:List[Any]		       =						'''mega'''


  def __init__(					self ,						snake_case_=3_0_5_2_2 ,						snake_case_=1_2_8 ,						snake_case_=4 ,						snake_case_=2_5_6 ,						snake_case_=1_6 ,						snake_case_=True ,						snake_case_=6_4 ,						snake_case_=False ,						snake_case_=-1 ,						snake_case_=None ,						snake_case_=True ,						snake_case_="scalenorm" ,						snake_case_=True ,						snake_case_="silu" ,						snake_case_="softmax" ,						snake_case_=0.1 ,						snake_case_=0.1 ,						snake_case_=0.1 ,						snake_case_=False ,						snake_case_=True ,						snake_case_=2_5_6 ,						snake_case_=True ,						snake_case_=0.1 ,						snake_case_=2_0_4_8 ,						snake_case_=False ,						snake_case_=2 ,						snake_case_=0.02 ,						snake_case_=0.2 ,						snake_case_=0.02 ,						snake_case_=1.0 ,						snake_case_=1 ,						snake_case_=0 ,						snake_case_=2 ,						snake_case_="rotary" ,						snake_case_=None ,						snake_case_=True ,						snake_case_=True ,						**snake_case_ ,						):




       '''simple docstring'''

       super().__init__(pad_token_id=snake_case_ ,						bos_token_id=snake_case_ ,						eos_token_id=snake_case_ ,						**snake_case_		)

       UpperCAmelCase_				:  Optional[Any]  = vocab_size
       UpperCAmelCase_				:  str  = hidden_size
       UpperCAmelCase_				:  List[str]  = num_hidden_layers
       UpperCAmelCase_				:  Tuple  = activation
       UpperCAmelCase_				:  List[Any]  = attention_activation
       UpperCAmelCase_				:  Dict  = intermediate_size
       UpperCAmelCase_				:  List[Any]  = ema_projection_size
       UpperCAmelCase_				:  int  = bidirectional
       UpperCAmelCase_				:  int  = shared_representation_size
       UpperCAmelCase_				:  Tuple  = use_chunking
       UpperCAmelCase_				:  int  = chunk_size
       UpperCAmelCase_				:  str  = truncation
       UpperCAmelCase_				:  Optional[int]  = normalize_before_mega
       UpperCAmelCase_				:  Union[str, Any]  = normalization_type
       UpperCAmelCase_				:  Dict  = norm_affine
       UpperCAmelCase_				:  Dict  = dropout_prob
       UpperCAmelCase_				:  Optional[int]  = hidden_dropout_prob
       UpperCAmelCase_				:  Tuple  = attention_probs_dropout_prob
       UpperCAmelCase_				:  Union[str, Any]  = use_feature_dropout
       UpperCAmelCase_				:  int  = use_normalized_ffn
       UpperCAmelCase_				:  Dict  = nffn_hidden_size
       UpperCAmelCase_				:  Optional[Any]  = normalize_before_ffn
       UpperCAmelCase_				:  Union[str, Any]  = nffn_activation_dropout_prob
       UpperCAmelCase_				:  Dict  = max_positions
       UpperCAmelCase_				:  Any  = add_token_type_embeddings
       UpperCAmelCase_				:  Dict  = type_vocab_size
       UpperCAmelCase_				:  Dict  = initializer_range
       UpperCAmelCase_				:  Dict  = ema_delta_alpha_range
       UpperCAmelCase_				:  str  = ema_beta_range
       UpperCAmelCase_				:  Optional[Any]  = ema_gamma_omega_range
       UpperCAmelCase_				:  str  = relative_positional_bias
       UpperCAmelCase_				:  str  = use_cache
       UpperCAmelCase_				:  List[str]  = classifier_dropout
       UpperCAmelCase_				:  List[Any]  = add_lm_hidden_dense_layer
       UpperCAmelCase_				:  Any  = 1  # not used but required by Hugging Face







class 	__SCREAMING_SNAKE_CASE		(			lowerCamelCase_			):







  '''simple docstring'''






  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       if self.task == "multiple-choice":
            UpperCAmelCase_				:  int  = {0: 'batch', 1: 'choice', 2: 'sequence'}
       else:
            UpperCAmelCase_				:  List[Any]  = {0: 'batch', 1: 'sequence'}
       return OrderedDict(
           [
               ('input_ids', dynamic_axis),
               ('attention_mask', dynamic_axis),
           ]		)







'''simple docstring'''


import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

import pandas as pd
import pyarrow as pa

import datasets
import datasets.config
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


if TYPE_CHECKING:
 import sqlitea

 import sqlalchemy


snake_case__	:     int      =    datasets.utils.logging.get_logger(__name__)






@dataclass
class 	__SCREAMING_SNAKE_CASE		(			datasets.BuilderConfig			):







  '''simple docstring'''






  lowerCamelCase_						:Union[str, "sqlalchemy.sql.Selectable"]		       =						None
  lowerCamelCase_						:Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]		       =						None
  lowerCamelCase_						:Optional[Union[str, List[str]]]		       =						None
  lowerCamelCase_						:bool		       =						True
  lowerCamelCase_						:Optional[Union[List, Tuple, Dict]]		       =						None
  lowerCamelCase_						:Optional[Union[List, Dict]]		       =						None
  lowerCamelCase_						:Optional[List[str]]		       =						None
  lowerCamelCase_						:Optional[int]		       =						1_0000
  lowerCamelCase_						:Optional[datasets.Features]		       =						None


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       if self.sql is None:
            raise ValueError('sql must be specified'		)
       if self.con is None:
            raise ValueError('con must be specified'		)


  def _UpperCamelCase	(					self ,						snake_case_ ,						snake_case_ = None ,						):




       '''simple docstring'''

       UpperCAmelCase_				:  List[Any]  = config_kwargs.copy()
       # We need to stringify the Selectable object to make its hash deterministic

       # The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
       UpperCAmelCase_				:  str  = config_kwargs['sql']
       if not isinstance(snake_case_ ,						snake_case_		):
            if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
                 import sqlalchemy

                 if isinstance(snake_case_ ,						sqlalchemy.sql.Selectable		):
                      UpperCAmelCase_				:  Any  = sqlalchemy.create_engine(config_kwargs['con'].split('://'		)[0] + '://'		)
                      UpperCAmelCase_				:  int  = str(sql.compile(dialect=engine.dialect		)		)
                      UpperCAmelCase_				:  str  = sql_str
                 else:
                      raise TypeError(
                          F'''Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(snake_case_		)}: {sql}'''		)
            else:
                 raise TypeError(
                     F'''Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(snake_case_		)}: {sql}'''		)
       UpperCAmelCase_				:  Union[str, Any]  = config_kwargs['con']
       if not isinstance(snake_case_ ,						snake_case_		):
            UpperCAmelCase_				:  Optional[Any]  = id(snake_case_		)
            logger.info(
                F'''SQL connection \'con\' of type {type(snake_case_		)} couldn\'t be hashed properly. To enable hashing, specify \'con\' as URI string instead.'''		)

       return super().create_config_id(snake_case_ ,						custom_features=snake_case_		)






  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  List[Any]  = {
           'index_col': self.index_col,
           'columns': self.columns,
           'params': self.params,
           'coerce_float': self.coerce_float,
           'parse_dates': self.parse_dates,
       }
       return pd_read_sql_kwargs







class 	__SCREAMING_SNAKE_CASE		(			datasets.ArrowBasedBuilder			):







  '''simple docstring'''






  lowerCamelCase_						:Any		       =						SqlConfig


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

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


  def _UpperCamelCase	(					self ,						snake_case_		):




       '''simple docstring'''

       return [datasets.SplitGenerator(name=datasets.Split.TRAIN ,						gen_kwargs={}		)]


  def _UpperCamelCase	(					self ,						snake_case_		):




       '''simple docstring'''

       if self.config.features is not None:
            UpperCAmelCase_				:  Tuple  = self.config.features.arrow_schema
            if all(not require_storage_cast(snake_case_		) for feature in self.config.features.values()		):
                 # cheaper cast
                 UpperCAmelCase_				:  Tuple  = pa.Table.from_arrays([pa_table[field.name] for field in schema] ,						schema=snake_case_		)
            else:
                 # more expensive cast; allows str <-> int/float or str to Audio for example
                 UpperCAmelCase_				:  Union[str, Any]  = table_cast(snake_case_ ,						snake_case_		)
       return pa_table






  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Optional[Any]  = self.config.chunksize
       UpperCAmelCase_				:  Optional[Any]  = pd.read_sql(
           self.config.sql ,						self.config.con ,						chunksize=snake_case_ ,						**self.config.pd_read_sql_kwargs		)
       UpperCAmelCase_				:  int  = [sql_reader] if chunksize is None else sql_reader
       for chunk_idx, df in enumerate(snake_case_		):
            UpperCAmelCase_				:  Optional[Any]  = pa.Table.from_pandas(snake_case_		)
            yield chunk_idx, self._cast_table(snake_case_		)






