



"""simple docstring"""


import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    NystromformerForMaskedLM,
					    NystromformerForMultipleChoice,
					    NystromformerForQuestionAnswering,
					    NystromformerForSequenceClassification,
					    NystromformerForTokenClassification,
					    NystromformerModel,
					)
					from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST
class       _a			:
			"""simple docstring"""



			def __init__(		self  :     List[Any]					, __UpperCamelCase  :     Any					, __UpperCamelCase  :     Tuple=1_3					, __UpperCamelCase  :     str=7					, __UpperCamelCase  :     str=True					, __UpperCamelCase  :     int=True					, __UpperCamelCase  :     int=True					, __UpperCamelCase  :     str=True					, __UpperCamelCase  :     Any=9_9					, __UpperCamelCase  :     Optional[int]=3_2					, __UpperCamelCase  :     Optional[int]=5					, __UpperCamelCase  :     Optional[int]=4					, __UpperCamelCase  :     Union[str, Any]=3_7					, __UpperCamelCase  :     List[str]="gelu"					, __UpperCamelCase  :     Any=0.1					, __UpperCamelCase  :     int=0.1					, __UpperCamelCase  :     Any=5_1_2					, __UpperCamelCase  :     List[str]=1_6					, __UpperCamelCase  :     Optional[Any]=2					, __UpperCamelCase  :     Optional[Any]=0.0_2					, __UpperCamelCase  :     List[Any]=3					, __UpperCamelCase  :     int=4					, __UpperCamelCase  :     Union[str, Any]=None					, )->Union[str, Any]:
							_UpperCAmelCase      		=     parent
							_UpperCAmelCase      		=     batch_size
							_UpperCAmelCase      		=     seq_length
							_UpperCAmelCase      		=     is_training
							_UpperCAmelCase      		=     use_input_mask
							_UpperCAmelCase      		=     use_token_type_ids
							_UpperCAmelCase      		=     use_labels
							_UpperCAmelCase      		=     vocab_size
							_UpperCAmelCase      		=     hidden_size
							_UpperCAmelCase      		=     num_hidden_layers
							_UpperCAmelCase      		=     num_attention_heads
							_UpperCAmelCase      		=     intermediate_size
							_UpperCAmelCase      		=     hidden_act
							_UpperCAmelCase      		=     hidden_dropout_prob
							_UpperCAmelCase      		=     attention_probs_dropout_prob
							_UpperCAmelCase      		=     max_position_embeddings
							_UpperCAmelCase      		=     type_vocab_size
							_UpperCAmelCase      		=     type_sequence_label_size
							_UpperCAmelCase      		=     initializer_range
							_UpperCAmelCase      		=     num_labels
							_UpperCAmelCase      		=     num_choices
							_UpperCAmelCase      		=     scope



			def 			lowercase__    (		self  :     Dict )->Tuple:
							_UpperCAmelCase      		=     ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size )

							_UpperCAmelCase      		=     None
							if self.use_input_mask:
											_UpperCAmelCase      		=     random_attention_mask([self.batch_size, self.seq_length] )

							_UpperCAmelCase      		=     None
							if self.use_token_type_ids:
											_UpperCAmelCase      		=     ids_tensor([self.batch_size, self.seq_length]					, self.type_vocab_size )

							_UpperCAmelCase      		=     None
							_UpperCAmelCase      		=     None
							_UpperCAmelCase      		=     None
							if self.use_labels:
											_UpperCAmelCase      		=     ids_tensor([self.batch_size]					, self.type_sequence_label_size )
											_UpperCAmelCase      		=     ids_tensor([self.batch_size, self.seq_length]					, self.num_labels )
											_UpperCAmelCase      		=     ids_tensor([self.batch_size]					, self.num_choices )

							_UpperCAmelCase      		=     self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



			def 			lowercase__    (		self  :     int )->List[str]:
							return NystromformerConfig(
							    vocab_size=self.vocab_size					, hidden_size=self.hidden_size					, num_hidden_layers=self.num_hidden_layers					, num_attention_heads=self.num_attention_heads					, intermediate_size=self.intermediate_size					, hidden_act=self.hidden_act					, hidden_dropout_prob=self.hidden_dropout_prob					, attention_probs_dropout_prob=self.attention_probs_dropout_prob					, max_position_embeddings=self.max_position_embeddings					, type_vocab_size=self.type_vocab_size					, is_decoder=snake_case_					, initializer_range=self.initializer_range					, )



			def 			lowercase__    (		self  :     Union[str, Any]					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     Optional[int]					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     Dict					, __UpperCamelCase  :     int					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     List[str] )->Optional[Any]:
							_UpperCAmelCase      		=     NystromformerModel(config=snake_case_ )
							model.to(snake_case_ )
							model.eval()
							_UpperCAmelCase      		=     model(snake_case_					, attention_mask=snake_case_					, token_type_ids=snake_case_ )
							_UpperCAmelCase      		=     model(snake_case_					, token_type_ids=snake_case_ )
							_UpperCAmelCase      		=     model(snake_case_ )
							self.parent.assertEqual(result.last_hidden_state.shape					, (self.batch_size, self.seq_length, self.hidden_size) )



			def 			lowercase__    (		self  :     List[str]					, __UpperCamelCase  :     int					, __UpperCamelCase  :     Dict					, __UpperCamelCase  :     Optional[int]					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     Union[str, Any]					, __UpperCamelCase  :     List[str]					, __UpperCamelCase  :     str )->Optional[int]:
							_UpperCAmelCase      		=     NystromformerForMaskedLM(config=snake_case_ )
							model.to(snake_case_ )
							model.eval()
							_UpperCAmelCase      		=     model(snake_case_					, attention_mask=snake_case_					, token_type_ids=snake_case_					, labels=snake_case_ )
							self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.vocab_size) )



			def 			lowercase__    (		self  :     Tuple					, __UpperCamelCase  :     Dict					, __UpperCamelCase  :     Any					, __UpperCamelCase  :     List[str]					, __UpperCamelCase  :     Optional[int]					, __UpperCamelCase  :     str					, __UpperCamelCase  :     Optional[int]					, __UpperCamelCase  :     int )->int:
							_UpperCAmelCase      		=     NystromformerForQuestionAnswering(config=snake_case_ )
							model.to(snake_case_ )
							model.eval()
							_UpperCAmelCase      		=     model(
							    snake_case_					, attention_mask=snake_case_					, token_type_ids=snake_case_					, start_positions=snake_case_					, end_positions=snake_case_					, )
							self.parent.assertEqual(result.start_logits.shape					, (self.batch_size, self.seq_length) )
							self.parent.assertEqual(result.end_logits.shape					, (self.batch_size, self.seq_length) )



			def 			lowercase__    (		self  :     List[Any]					, __UpperCamelCase  :     int					, __UpperCamelCase  :     Union[str, Any]					, __UpperCamelCase  :     Dict					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     List[Any]					, __UpperCamelCase  :     Tuple )->Union[str, Any]:
							_UpperCAmelCase      		=     self.num_labels
							_UpperCAmelCase      		=     NystromformerForSequenceClassification(snake_case_ )
							model.to(snake_case_ )
							model.eval()
							_UpperCAmelCase      		=     model(snake_case_					, attention_mask=snake_case_					, token_type_ids=snake_case_					, labels=snake_case_ )
							self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_labels) )



			def 			lowercase__    (		self  :     str					, __UpperCamelCase  :     Optional[int]					, __UpperCamelCase  :     Optional[int]					, __UpperCamelCase  :     Any					, __UpperCamelCase  :     str					, __UpperCamelCase  :     Optional[Any]					, __UpperCamelCase  :     Tuple					, __UpperCamelCase  :     str )->List[str]:
							_UpperCAmelCase      		=     self.num_labels
							_UpperCAmelCase      		=     NystromformerForTokenClassification(config=snake_case_ )
							model.to(snake_case_ )
							model.eval()
							_UpperCAmelCase      		=     model(snake_case_					, attention_mask=snake_case_					, token_type_ids=snake_case_					, labels=snake_case_ )
							self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.num_labels) )



			def 			lowercase__    (		self  :     Tuple					, __UpperCamelCase  :     Any					, __UpperCamelCase  :     Tuple					, __UpperCamelCase  :     Union[str, Any]					, __UpperCamelCase  :     List[Any]					, __UpperCamelCase  :     Any					, __UpperCamelCase  :     Any					, __UpperCamelCase  :     List[Any] )->Any:
							_UpperCAmelCase      		=     self.num_choices
							_UpperCAmelCase      		=     NystromformerForMultipleChoice(config=snake_case_ )
							model.to(snake_case_ )
							model.eval()
							_UpperCAmelCase      		=     input_ids.unsqueeze(1 ).expand(-1					, self.num_choices					, -1 ).contiguous()
							_UpperCAmelCase      		=     token_type_ids.unsqueeze(1 ).expand(-1					, self.num_choices					, -1 ).contiguous()
							_UpperCAmelCase      		=     input_mask.unsqueeze(1 ).expand(-1					, self.num_choices					, -1 ).contiguous()
							_UpperCAmelCase      		=     model(
							    snake_case_					, attention_mask=snake_case_					, token_type_ids=snake_case_					, labels=snake_case_					, )
							self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_choices) )



			def 			lowercase__    (		self  :     List[Any] )->List[str]:
							_UpperCAmelCase      		=     self.prepare_config_and_inputs()
							(
							    (
							    _UpperCAmelCase			
							)		,    (
							    _UpperCAmelCase			
							)		,    (
							    _UpperCAmelCase			
							)		,    (
							    _UpperCAmelCase			
							)		,    (
							    _UpperCAmelCase			
							)		,    (
							    _UpperCAmelCase			
							)		,    (
							    _UpperCAmelCase			
							)		,    			
							)      		=     config_and_inputs
							_UpperCAmelCase      		=     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
							return config, inputs_dict



@require_torch
class       _a			(       __snake_case      ,				__snake_case      ,				unittest.TestCase):
			"""simple docstring"""
			UpperCamelCase__         =						(
			    (
			        NystromformerModel,
			        NystromformerForMaskedLM,
			        NystromformerForMultipleChoice,
			        NystromformerForQuestionAnswering,
			        NystromformerForSequenceClassification,
			        NystromformerForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			UpperCamelCase__         =						(
			    {
			        """feature-extraction""": NystromformerModel,
			        """fill-mask""": NystromformerForMaskedLM,
			        """question-answering""": NystromformerForQuestionAnswering,
			        """text-classification""": NystromformerForSequenceClassification,
			        """token-classification""": NystromformerForTokenClassification,
			        """zero-shot""": NystromformerForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			UpperCamelCase__         =						False
			UpperCamelCase__         =						False



			def 			lowercase__    (		self  :     int )->Tuple:
							_UpperCAmelCase      		=     NystromformerModelTester(self )
							_UpperCAmelCase      		=     ConfigTester(self					, config_class=snake_case_					, hidden_size=3_7 )



			def 			lowercase__    (		self  :     List[Any] )->Dict:
							self.config_tester.run_common_tests()



			def 			lowercase__    (		self  :     List[str] )->Optional[Any]:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case_ )



			def 			lowercase__    (		self  :     List[Any] )->List[str]:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
											_UpperCAmelCase      		=     type
											self.model_tester.create_and_check_model(*snake_case_ )



			def 			lowercase__    (		self  :     Optional[Any] )->int:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*snake_case_ )



			def 			lowercase__    (		self  :     int )->Dict:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*snake_case_ )



			def 			lowercase__    (		self  :     str )->Any:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*snake_case_ )



			def 			lowercase__    (		self  :     Any )->Dict:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*snake_case_ )



			def 			lowercase__    (		self  :     int )->Optional[int]:
							_UpperCAmelCase      		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*snake_case_ )



			@slow
			def 			lowercase__    (		self  :     int )->Optional[Any]:
							for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_UpperCAmelCase      		=     NystromformerModel.from_pretrained(snake_case_ )
											self.assertIsNotNone(snake_case_ )



@require_torch
class       _a			(       unittest.TestCase):
			"""simple docstring"""



			@slow
			def 			lowercase__    (		self  :     str )->Union[str, Any]:
							_UpperCAmelCase      		=     NystromformerModel.from_pretrained('''uw-madison/nystromformer-512''' )
							_UpperCAmelCase      		=     torch.tensor([[0, 1, 2, 3, 4, 5]] )

							with torch.no_grad():
											_UpperCAmelCase      		=     model(snake_case_ )[0]

							_UpperCAmelCase      		=     torch.Size((1, 6, 7_6_8) )
							self.assertEqual(output.shape					, snake_case_ )

							_UpperCAmelCase      		=     torch.tensor(
							    [[[-0.4_5_3_2, -0.0_9_3_6, 0.5_1_3_7], [-0.2_6_7_6, 0.0_6_2_8, 0.6_1_8_6], [-0.3_6_2_9, -0.1_7_2_6, 0.4_7_1_6]]] )

							self.assertTrue(torch.allclose(output[:, :3, :3]					, snake_case_					, atol=1e-4 ) )



			@slow
			def 			lowercase__    (		self  :     Tuple )->Tuple:
							_UpperCAmelCase      		=     '''the [MASK] of Belgium is Brussels'''

							_UpperCAmelCase      		=     AutoTokenizer.from_pretrained('''uw-madison/nystromformer-512''' )
							_UpperCAmelCase      		=     NystromformerForMaskedLM.from_pretrained('''uw-madison/nystromformer-512''' )

							_UpperCAmelCase      		=     tokenizer(snake_case_					, return_tensors='''pt''' )

							with torch.no_grad():
											_UpperCAmelCase      		=     model(encoding.input_ids ).logits

							_UpperCAmelCase      		=     token_logits[:, 2, :].argmax(-1 )[0]

							self.assertEqual(tokenizer.decode(snake_case_ )					, '''capital''' )










# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import subprocess

from packaging.version import Version, parse

from accelerate.commands.config.config_args import default_config_file, load_config_from_file


__A							:			Optional[int]         =			"Run commands across TPU VMs for initial setup before running `accelerate launch`."





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE=None				)							-> str:




    """simple docstring"""




    if subparsers is not None:
        _A     =							subparsers.add_parser('tpu-config'							,     description=_description				)
    else:
        _A     =							argparse.ArgumentParser('Accelerate tpu-config command'							,     description=_description				)
    # Core arguments
    _A     =							parser.add_argument_group(
        'Config Arguments'							,     'Arguments that can be configured through `accelerate config`.'				)
    config_args.add_argument(
        '--config_file'							,     type=_SCREAMING_SNAKE_CASE							,     default=_SCREAMING_SNAKE_CASE							,     help='Path to the config file to use for accelerate.'							,     )
    config_args.add_argument(
        '--tpu_name'							,     default=_SCREAMING_SNAKE_CASE							,     help='The name of the TPU to use. If not specified, will use the TPU specified in the config file.'							,     )
    config_args.add_argument(
        '--tpu_zone'							,     default=_SCREAMING_SNAKE_CASE							,     help='The zone of the TPU to use. If not specified, will use the zone specified in the config file.'							,     )
    _A     =							parser.add_argument_group('TPU Arguments'							,     'Arguments for options ran inside the TPU.'				)
    pod_args.add_argument(
        '--use_alpha'							,     action='store_true'							,     help='Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`.'							,     )
    pod_args.add_argument(
        '--command_file'							,     default=_SCREAMING_SNAKE_CASE							,     help='The path to the file containing the commands to run on the pod on startup.'							,     )
    pod_args.add_argument(
        '--command'							,     action='append'							,     nargs='+'							,     help='A command to run on the pod. Can be passed multiple times.'							,     )
    pod_args.add_argument(
        '--install_accelerate'							,     action='store_true'							,     help='Whether to install accelerate on the pod. Defaults to False.'							,     )
    pod_args.add_argument(
        '--accelerate_version'							,     default='latest'							,     help='The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify \'dev\' to install from GitHub.'							,     )
    pod_args.add_argument(
        '--debug'							,     action='store_true'							,     help='If set, will print the command that would be run instead of running it.'				)

    if subparsers is not None:
        parser.set_defaults(func=_SCREAMING_SNAKE_CASE				)
    return parser





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE				)							-> int:




    """simple docstring"""




    _A     =							None

    # Get the default from the config file if it exists.
    if args.config_file is not None or os.path.isfile(_SCREAMING_SNAKE_CASE				):
        _A     =							load_config_from_file(args.config_file				)
        if not args.command_file and defaults.command_file is not None and not args.command:
            _A     =							defaults.command_file
        if not args.command and defaults.commands is not None:
            _A     =							defaults.commands
        if not args.tpu_name:
            _A     =							defaults.tpu_name
        if not args.tpu_zone:
            _A     =							defaults.tpu_zone
    if args.accelerate_version == "dev":
        _A     =							'git+https://github.com/huggingface/accelerate.git'
    elif args.accelerate_version == "latest":
        _A     =							'accelerate -U'
    elif isinstance(parse(args.accelerate_version				)							,     _SCREAMING_SNAKE_CASE				):
        _A     =							F"accelerate=={args.accelerate_version}"

    if not args.command_file and not args.command:
        raise ValueError('You must specify either a command file or a command to run on the pod.'				)

    if args.command_file:
        with open(args.command_file							,     'r'				) as f:
            _A     =							[f.read().splitlines()]

    # To turn list of lists into list of strings
    if isinstance(args.command[0]							,     _SCREAMING_SNAKE_CASE				):
        _A     =							[line for cmd in args.command for line in cmd]
    # Default to the shared folder and install accelerate
    _A     =							['cd /usr/share']
    if args.install_accelerate:
        new_cmd += [F"pip install {args.accelerate_version}"]
    new_cmd += args.command
    _A     =							'; '.join(_SCREAMING_SNAKE_CASE				)

    # Then send it to gcloud
    # Eventually try to use google-api-core to do this instead of subprocess
    _A     =							['gcloud']
    if args.use_alpha:
        cmd += ["alpha"]
    cmd += [
        "compute",
        "tpus",
        "tpu-vm",
        "ssh",
        args.tpu_name,
        "--zone",
        args.tpu_zone,
        "--command",
        args.command,
        "--worker",
        "all",
    ]
    if args.debug:
        print(F"Running {' '.join(_SCREAMING_SNAKE_CASE				)}"				)
        return
    subprocess.run(_SCREAMING_SNAKE_CASE				)
    print('Successfully setup pod.'				)





def 				__lowerCAmelCase(       )							-> Tuple:




    """simple docstring"""




    _A     =							tpu_command_parser()
    _A     =							parser.parse_args()

    tpu_command_launcher(_SCREAMING_SNAKE_CASE				)





