




"""simple docstring"""





from ...configuration_utils import PretrainedConfig
from ...utils import logging


A   =       logging.get_logger(__name__)

A   =       {
    """google/switch-base-8""": """https://huggingface.co/google/switch-base-8/blob/main/config.json""",
}



class      a__				(	__magic_name__   ):
			lowercase_    =	"switch_transformers"
			lowercase_    =	["past_key_values"]
			lowercase_    =	{"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}


			def __init__(      self      :    str      ,  UpperCamelCase_      :    int=32128      ,  UpperCamelCase_      :    Optional[Any]=768      ,  UpperCamelCase_      :    Optional[Any]=64      ,  UpperCamelCase_      :    Optional[int]=2048      ,  UpperCamelCase_      :    List[str]=64      ,  UpperCamelCase_      :    int=12      ,  UpperCamelCase_      :    List[Any]=3      ,  UpperCamelCase_      :    Union[str, Any]=12      ,  UpperCamelCase_      :    Optional[int]=3      ,  UpperCamelCase_      :    str=12      ,  UpperCamelCase_      :    Optional[int]=8      ,  UpperCamelCase_      :    List[str]=False      ,  UpperCamelCase_      :    Any=0.01      ,  UpperCamelCase_      :    Union[str, Any]="float32"      ,  UpperCamelCase_      :    List[str]=False      ,  UpperCamelCase_      :    int=32      ,  UpperCamelCase_      :    Tuple=128      ,  UpperCamelCase_      :    List[str]=0.1      ,  UpperCamelCase_      :    Union[str, Any]=1e-6      ,  UpperCamelCase_      :    Union[str, Any]=0.001      ,  UpperCamelCase_      :    Any=0.001      ,  UpperCamelCase_      :    Dict=1.0      ,  UpperCamelCase_      :    Optional[int]="relu"      ,  UpperCamelCase_      :    List[str]=True      ,  UpperCamelCase_      :    List[str]=False      ,  UpperCamelCase_      :    int=True      ,  UpperCamelCase_      :    Tuple=0      ,  UpperCamelCase_      :    Optional[Any]=1      ,  **UpperCamelCase_      :    Optional[int]      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					vocab_size
						__UpperCAmelCase	: Optional[Any]		=					d_model
						__UpperCAmelCase	: Union[str, Any]		=					d_kv
						__UpperCAmelCase	: str		=					d_ff

						__UpperCAmelCase	: Dict		=					num_sparse_encoder_layers

						__UpperCAmelCase	: Dict		=					num_layers
						__UpperCAmelCase	: str		=					(
						    num_decoder_layers if num_decoder_layers is not None else self.num_layers
						)  # default = symmetry
						__UpperCAmelCase	: int		=					num_sparse_decoder_layers

						# This tells us, each how many encoder layer we'll have to set a sparse layer.
						if self.num_sparse_encoder_layers > 0:
									__UpperCAmelCase	: Any		=					self.num_layers // self.num_sparse_encoder_layers
						else:
									__UpperCAmelCase	: Optional[Any]		=					self.num_layers  # HACK: this will create 0 sparse layers

						# This tells us, each how many encoder layer we'll have to set a sparse layer.
						if self.num_sparse_decoder_layers > 0:
									__UpperCAmelCase	: Union[str, Any]		=					self.num_decoder_layers // self.num_sparse_decoder_layers
						else:
									__UpperCAmelCase	: Tuple		=					self.num_decoder_layers  # HACK: this will create 0 sparse layers

						__UpperCAmelCase	: List[Any]		=					num_heads
						__UpperCAmelCase	: Union[str, Any]		=					num_experts
						__UpperCAmelCase	: str		=					expert_capacity
						__UpperCAmelCase	: Any		=					router_bias
						__UpperCAmelCase	: Any		=					router_jitter_noise
						if router_dtype not in ["float32", "float16", "bfloat16"]:
									raise ValueError(F"`router_dtype` must be one of \'float32\', \'float16\' or \'bfloat16\', got {router_dtype}")
						__UpperCAmelCase	: List[str]		=					router_dtype

						__UpperCAmelCase	: Any		=					router_ignore_padding_tokens
						__UpperCAmelCase	: str		=					relative_attention_num_buckets
						__UpperCAmelCase	: Optional[int]		=					relative_attention_max_distance

						__UpperCAmelCase	: Optional[Any]		=					dropout_rate
						__UpperCAmelCase	: Optional[int]		=					layer_norm_epsilon
						__UpperCAmelCase	: Optional[Any]		=					initializer_factor
						__UpperCAmelCase	: List[Any]		=					feed_forward_proj
						__UpperCAmelCase	: List[Any]		=					use_cache
						__UpperCAmelCase	: int		=					add_router_probs

						__UpperCAmelCase	: Tuple		=					router_z_loss_coef
						__UpperCAmelCase	: Tuple		=					router_aux_loss_coef

						__UpperCAmelCase	: str		=					self.feed_forward_proj.split("-")
						__UpperCAmelCase	: Union[str, Any]		=					act_info[-1]
						__UpperCAmelCase	: Tuple		=					act_info[0] == "gated"

						if len(__UpperCamelCase) > 1 and act_info[0] != "gated" or len(__UpperCamelCase) > 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'")

						# for backwards compatibility
						if feed_forward_proj == "gated-gelu":
									__UpperCAmelCase	: Union[str, Any]		=					"gelu_new"

						super().__init__(
						    pad_token_id=__UpperCamelCase      ,  eos_token_id=__UpperCamelCase      ,  is_encoder_decoder=__UpperCamelCase      ,  **__UpperCamelCase      ,  )








"""simple docstring"""





from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader



class      a__				(	__magic_name__   ):


			def __init__(      self      :    List[Any]      ,  UpperCamelCase_      :    pyspark.sql.DataFrame      ,  UpperCamelCase_      :    Optional[NamedSplit] = None      ,  UpperCamelCase_      :    Optional[Features] = None      ,  UpperCamelCase_      :    bool = True      ,  UpperCamelCase_      :    str = None      ,  UpperCamelCase_      :    bool = False      ,  UpperCamelCase_      :    str = None      ,  UpperCamelCase_      :    bool = True      ,  UpperCamelCase_      :    str = "arrow"      ,  **UpperCamelCase_      :    Tuple      ,  ):


						"""simple docstring"""




						super().__init__(
						    split=UpperCamelCase_      ,  features=UpperCamelCase_      ,  cache_dir=UpperCamelCase_      ,  keep_in_memory=UpperCamelCase_      ,  streaming=UpperCamelCase_      ,  **UpperCamelCase_      ,  )
						__UpperCAmelCase	: List[Any]		=					load_from_cache_file
						__UpperCAmelCase	: List[Any]		=					file_format
						__UpperCAmelCase	: str		=					Spark(
						    df=UpperCamelCase_      ,  features=UpperCamelCase_      ,  cache_dir=UpperCamelCase_      ,  working_dir=UpperCamelCase_      ,  **UpperCamelCase_      ,  )


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


						"""simple docstring"""




						if self.streaming:
									return self.builder.as_streaming_dataset(split=self.split)
						__UpperCAmelCase	: List[Any]		=					None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
						self.builder.download_and_prepare(
						    download_mode=UpperCamelCase_      ,  file_format=self._file_format      ,  )
						return self.builder.as_dataset(split=self.split)


