



from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
       import tensorflow as tf


lowercase_						=       logging.get_logger(__name__)
@dataclass
class 			_snake_case     (   snake_case_):
    UpperCamelCase__			:   Dict       						=[
        """no_inference""",
        """no_cuda""",
        """no_tpu""",
        """no_speed""",
        """no_memory""",
        """no_env_print""",
        """no_multi_process""",
    ]




    def __init__(  self   :  Union[str, Any],   **__lowercase   :  Optional[int]  ):
       for deprecated_arg in self.deprecated_args:
          if deprecated_arg in kwargs:
             lowercase__   					=			deprecated_arg[3:]
             lowercase__   					=			not kwargs.pop(__lowercase  )
             logger.warning(
                 F'''{deprecated_arg} is depreciated. Please use --no-{positive_arg} or'''
                 F''' {positive_arg}={kwargs[positive_arg]}'''  )
       lowercase__   					=			kwargs.pop("tpu_name",   self.tpu_name  )
       lowercase__   					=			kwargs.pop("device_idx",   self.device_idx  )
       lowercase__   					=			kwargs.pop("eager_mode",   self.eager_mode  )
       lowercase__   					=			kwargs.pop("use_xla",   self.use_xla  )
       super().__init__(**__lowercase  )

    UpperCamelCase__			:   Tuple       						=field(
        default=snake_case_      ,		metadata={"""help""": """Name of TPU"""}      ,		)
    UpperCamelCase__			:   List[str]       						=field(
        default=0      ,		metadata={"""help""": """CPU / GPU device index. Defaults to 0."""}      ,		)
    UpperCamelCase__			:   Optional[int]       						=field(default=snake_case_      ,		metadata={"""help""": """Benchmark models in eager model."""})
    UpperCamelCase__			:   Optional[Any]       						=field(
        default=snake_case_      ,		metadata={
            """help""": """Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`."""
        }      ,		)




    @cached_property
    def 							A__ (  self   :  Any  ):
       requires_backends(self,   ["tf"]  )
       lowercase__   					=			None
       if self.tpu:
          try:
             if self.tpu_name:
                lowercase__   					=			tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name  )
             else:
                lowercase__   					=			tf.distribute.cluster_resolver.TPUClusterResolver()
          except ValueError:
             lowercase__   					=			None
       return tpu




    @cached_property
    def 							A__ (  self   :  Union[str, Any]  ):
       requires_backends(self,   ["tf"]  )
       if self.is_tpu:
          tf.config.experimental_connect_to_cluster(self._setup_tpu  )
          tf.tpu.experimental.initialize_tpu_system(self._setup_tpu  )

          lowercase__   					=			tf.distribute.TPUStrategy(self._setup_tpu  )
       else:
          # currently no multi gpu is allowed
          if self.is_gpu:
             # TODO: Currently only single GPU is supported
             tf.config.set_visible_devices(self.gpu_list[self.device_idx],   "GPU"  )
             lowercase__   					=			tf.distribute.OneDeviceStrategy(device=F'''/gpu:{self.device_idx}'''  )
          else:
             tf.config.set_visible_devices([],   "GPU"  )  # disable GPU
             lowercase__   					=			tf.distribute.OneDeviceStrategy(device=F'''/cpu:{self.device_idx}'''  )

       return strategy




    @property
    def 							A__ (  self   :  str  ):
       requires_backends(self,   ["tf"]  )
       return self._setup_tpu is not None




    @property
    def 							A__ (  self   :  Any  ):
       requires_backends(self,   ["tf"]  )
       return self._setup_strategy




    @property
    def 							A__ (  self   :  int  ):
       requires_backends(self,   ["tf"]  )
       return tf.config.list_physical_devices("GPU"  )




    @property
    def 							A__ (  self   :  Tuple  ):
       requires_backends(self,   ["tf"]  )
       if self.cuda:
          return len(self.gpu_list  )
       return 0




    @property
    def 							A__ (  self   :  List[str]  ):
       return self.n_gpu > 0









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


lowercase_						=       logging.get_logger(__name__)

lowercase_						=       {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   Tuple       						="""nllb-moe"""
				UpperCamelCase__			:   Dict       						=["""past_key_values"""]
				UpperCamelCase__			:   List[Any]       						={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




				def __init__(  self   :  List[str],   __lowercase   :  str=12_8112,   __lowercase   :  List[Any]=1024,   __lowercase   :  int=12,   __lowercase   :  Union[str, Any]=4096,   __lowercase   :  Union[str, Any]=16,   __lowercase   :  Tuple=12,   __lowercase   :  int=4096,   __lowercase   :  Any=16,   __lowercase   :  List[str]=0.05,   __lowercase   :  Optional[int]=0.05,   __lowercase   :  int=True,   __lowercase   :  Optional[Any]=True,   __lowercase   :  Optional[Any]="relu",   __lowercase   :  str=1024,   __lowercase   :  Optional[Any]=0.1,   __lowercase   :  List[Any]=0.1,   __lowercase   :  Union[str, Any]=0.0,   __lowercase   :  Tuple=0.02,   __lowercase   :  Optional[Any]=2,   __lowercase   :  Optional[int]=True,   __lowercase   :  List[Any]=False,   __lowercase   :  Any="float32",   __lowercase   :  Union[str, Any]=False,   __lowercase   :  Optional[int]=128,   __lowercase   :  Union[str, Any]=64,   __lowercase   :  Optional[int]=4,   __lowercase   :  Dict=4,   __lowercase   :  Any=0.001,   __lowercase   :  Optional[int]=0.001,   __lowercase   :  int="all",   __lowercase   :  Any=False,   __lowercase   :  Optional[int]=False,   __lowercase   :  Dict=1.0,   __lowercase   :  int=0.2,   __lowercase   :  str=1,   __lowercase   :  Optional[Any]=0,   __lowercase   :  Optional[int]=2,   __lowercase   :  Any=False,   **__lowercase   :  List[str],   ):
							lowercase__   					=			vocab_size
							lowercase__   					=			max_position_embeddings
							lowercase__   					=			d_model
							lowercase__   					=			encoder_ffn_dim
							lowercase__   					=			encoder_layers
							lowercase__   					=			encoder_attention_heads
							lowercase__   					=			decoder_ffn_dim
							lowercase__   					=			decoder_layers
							lowercase__   					=			decoder_attention_heads
							lowercase__   					=			dropout
							lowercase__   					=			attention_dropout
							lowercase__   					=			activation_dropout
							lowercase__   					=			activation_function
							lowercase__   					=			init_std
							lowercase__   					=			encoder_layerdrop
							lowercase__   					=			decoder_layerdrop
							lowercase__   					=			use_cache
							lowercase__   					=			encoder_layers
							lowercase__   					=			scale_embedding  # scale factor will be sqrt(d_model) if True
							lowercase__   					=			router_z_loss_coef
							lowercase__   					=			router_aux_loss_coef
							lowercase__   					=			decoder_sparse_step
							lowercase__   					=			encoder_sparse_step
							lowercase__   					=			num_experts
							lowercase__   					=			expert_capacity
							lowercase__   					=			router_bias
							if router_dtype not in ["float32", "float16", "bfloat16"]:
										raise ValueError(F'''`router_dtype` must be one of \'float32\', \'float16\' or \'bfloat16\', got {router_dtype}'''  )
							lowercase__   					=			router_dtype

							lowercase__   					=			router_ignore_padding_tokens
							lowercase__   					=			batch_prioritized_routing
							lowercase__   					=			second_expert_policy
							lowercase__   					=			normalize_router_prob_before_dropping
							lowercase__   					=			moe_eval_capacity_token_fraction
							lowercase__   					=			moe_token_dropout
							lowercase__   					=			output_router_logits
							super().__init__(
							    pad_token_id=__lowercase,   bos_token_id=__lowercase,   eos_token_id=__lowercase,   is_encoder_decoder=__lowercase,   decoder_start_token_id=__lowercase,   **__lowercase,   )




