




import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

UpperCamelCase__					    =    logging.getLogger()







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Union[str, Any]  ):
					__lowerCAmelCase            = {}
					__lowerCAmelCase            = os.path.join(SCREAMING_SNAKE_CASE_   ,   "all_results.json"  )
					if os.path.exists(SCREAMING_SNAKE_CASE_  ):
										with open(SCREAMING_SNAKE_CASE_   ,   "r"  ) as f:
															__lowerCAmelCase            = json.load(SCREAMING_SNAKE_CASE_  )
					else:
										raise ValueError(F"""can\'t find {path}"""  )
					return results


UpperCamelCase__					    =    logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)



@require_torch_tpu
class     a__					(      __UpperCAmelCase   ):






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											import xla_spawn

											__lowerCAmelCase            = self.get_auto_remove_tmp_dir()
											__lowerCAmelCase            = f"""
            ./examples/pytorch/text-classification/run_glue.py
            --num_cores=8
            ./examples/pytorch/text-classification/run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --do_train
            --do_eval
            --debug tpu_metrics_debug
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --max_steps=10
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

											with patch.object(_A	,    "argv"	,    _A    ):
																__lowerCAmelCase            = time()
																xla_spawn.main()
																__lowerCAmelCase            = time()

																__lowerCAmelCase            = get_results(_A    )
																self.assertGreaterEqual(result["eval_accuracy"]	,    0.75    )

																# Assert that the script takes less than 500 seconds to make sure it doesn't hang.
																self.assertLess(end - start	,    5_0_0    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											import xla_spawn

											__lowerCAmelCase            = "\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            ".split()
											with patch.object(_A	,    "argv"	,    _A    ):
																xla_spawn.main()










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


UpperCamelCase__					    =    logging.get_logger(__name__)

UpperCamelCase__					    =    {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class     a__					(      snake_case__   ):
						_a							:					Tuple      =					"""nllb-moe"""
						_a							:					Dict      =					["""past_key_values"""]
						_a							:					Optional[int]      =					{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






						def __init__(     self	,    _A=1_2_8_1_1_2	,    _A=1_0_2_4	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=0.05	,    _A=0.05	,    _A=True	,    _A=True	,    _A="relu"	,    _A=1_0_2_4	,    _A=0.1	,    _A=0.1	,    _A=0.0	,    _A=0.02	,    _A=2	,    _A=True	,    _A=False	,    _A="float32"	,    _A=False	,    _A=1_2_8	,    _A=6_4	,    _A=4	,    _A=4	,    _A=0.0_01	,    _A=0.0_01	,    _A="all"	,    _A=False	,    _A=False	,    _A=1.0	,    _A=0.2	,    _A=1	,    _A=0	,    _A=2	,    _A=False	,    **_A	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = d_model
											__lowerCAmelCase            = encoder_ffn_dim
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = encoder_attention_heads
											__lowerCAmelCase            = decoder_ffn_dim
											__lowerCAmelCase            = decoder_layers
											__lowerCAmelCase            = decoder_attention_heads
											__lowerCAmelCase            = dropout
											__lowerCAmelCase            = attention_dropout
											__lowerCAmelCase            = activation_dropout
											__lowerCAmelCase            = activation_function
											__lowerCAmelCase            = init_std
											__lowerCAmelCase            = encoder_layerdrop
											__lowerCAmelCase            = decoder_layerdrop
											__lowerCAmelCase            = use_cache
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = scale_embedding  # scale factor will be sqrt(d_model) if True
											__lowerCAmelCase            = router_z_loss_coef
											__lowerCAmelCase            = router_aux_loss_coef
											__lowerCAmelCase            = decoder_sparse_step
											__lowerCAmelCase            = encoder_sparse_step
											__lowerCAmelCase            = num_experts
											__lowerCAmelCase            = expert_capacity
											__lowerCAmelCase            = 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}"""    )
											__lowerCAmelCase            = router_dtype

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




