





'''simple docstring'''




# 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


lowercase__       : Optional[Any]         =       '''Run commands across TPU VMs for initial setup before running `accelerate launch`.'''






def 			_lowerCAmelCase     (    __snake_case    :					str=None     )	->	int:
     if subparsers is not None:
          __A		:					Tuple						=			subparsers.add_parser('tpu-config'					,				description=_description     )
     else:
          __A		:					Dict						=			argparse.ArgumentParser('Accelerate tpu-config command'					,				description=_description     )
     # Core arguments
     __A		:					Dict						=			parser.add_argument_group(
         'Config Arguments'					,				'Arguments that can be configured through `accelerate config`.'     )
     config_args.add_argument(
         '--config_file'					,				type=__snake_case					,				default=__snake_case					,				help='Path to the config file to use for accelerate.'					,				)
     config_args.add_argument(
         '--tpu_name'					,				default=__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=__snake_case					,				help='The zone of the TPU to use. If not specified, will use the zone specified in the config file.'					,				)
     __A		:					Optional[Any]						=			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=__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=__snake_case     )
     return parser






def 			_lowerCAmelCase     (    __snake_case    :					Dict     )	->	Optional[int]:
     __A		:					Tuple						=			None

     # Get the default from the config file if it exists.
     if args.config_file is not None or os.path.isfile(__snake_case     ):
          __A		:					Tuple						=			load_config_from_file(args.config_file     )
          if not args.command_file and defaults.command_file is not None and not args.command:
               __A		:					Optional[int]						=			defaults.command_file
          if not args.command and defaults.commands is not None:
               __A		:					List[str]						=			defaults.commands
          if not args.tpu_name:
               __A		:					Tuple						=			defaults.tpu_name
          if not args.tpu_zone:
               __A		:					Union[str, Any]						=			defaults.tpu_zone
     if args.accelerate_version == "dev":
          __A		:					Union[str, Any]						=			'git+https://github.com/huggingface/accelerate.git'
     elif args.accelerate_version == "latest":
          __A		:					Optional[Any]						=			'accelerate -U'
     elif isinstance(parse(args.accelerate_version     )					,				__snake_case     ):
          __A		:					Tuple						=			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		:					str						=			[f.read().splitlines()]

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

     # Then send it to gcloud
     # Eventually try to use google-api-core to do this instead of subprocess
     __A		:					Dict						=			['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(__snake_case     )}'     )
          return
     subprocess.run(__snake_case     )
     print('Successfully setup pod.'     )







def 			_lowerCAmelCase     (    )	->	List[str]:
     __A		:					int						=			tpu_command_parser()
     __A		:					Union[str, Any]						=			parser.parse_args()

     tpu_command_launcher(__snake_case     )






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
    import torch

class   lowerCAmelCase				:
 def __init__( self    : List[str]     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : str=14     ,	UpperCAmelCase    : Any=7     ,	UpperCAmelCase    : Dict=True     ,	UpperCAmelCase    : List[str]=True     ,	UpperCAmelCase    : Any=False     ,	UpperCAmelCase    : Tuple=True     ,	UpperCAmelCase    : int=99     ,	UpperCAmelCase    : str=32     ,	UpperCAmelCase    : int=4     ,	UpperCAmelCase    : Optional[Any]=4     ,	UpperCAmelCase    : Any=4     ,	UpperCAmelCase    : Optional[int]=37     ,	UpperCAmelCase    : Union[str, Any]="gelu"     ,	UpperCAmelCase    : List[Any]=0.1     ,	UpperCAmelCase    : Dict=0.1     ,	UpperCAmelCase    : Optional[int]=512     ,	UpperCAmelCase    : List[Any]=0.0_2     ,	)     ->      List[Any]:
     lowerCamelCase__       : Tuple   					=						parent
     lowerCamelCase__       : Tuple   					=						batch_size
     lowerCamelCase__       : Optional[Any]   					=						seq_length
     lowerCamelCase__       : str   					=						is_training
     lowerCamelCase__       : Any   					=						use_input_mask
     lowerCamelCase__       : Optional[Any]   					=						use_token_type_ids
     lowerCamelCase__       : Optional[Any]   					=						use_labels
     lowerCamelCase__       : Optional[Any]   					=						vocab_size
     lowerCamelCase__       : Optional[int]   					=						hidden_size
     lowerCamelCase__       : Optional[Any]   					=						rotary_dim
     lowerCamelCase__       : Optional[Any]   					=						num_hidden_layers
     lowerCamelCase__       : Optional[Any]   					=						num_attention_heads
     lowerCamelCase__       : Tuple   					=						intermediate_size
     lowerCamelCase__       : Union[str, Any]   					=						hidden_act
     lowerCamelCase__       : List[Any]   					=						hidden_dropout_prob
     lowerCamelCase__       : int   					=						attention_probs_dropout_prob
     lowerCamelCase__       : Dict   					=						max_position_embeddings
     lowerCamelCase__       : Tuple   					=						initializer_range
     lowerCamelCase__       : Optional[int]   					=						None
     lowerCamelCase__       : int   					=						vocab_size - 1
     lowerCamelCase__       : str   					=						vocab_size - 1
     lowerCamelCase__       : str   					=						vocab_size - 1
 def 	A_      ( self    : str      )     ->      int:
     lowerCamelCase__       : Optional[int]   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size      )

     lowerCamelCase__       : Optional[Any]   					=						None
     if self.use_input_mask:
         lowerCamelCase__       : List[str]   					=						random_attention_mask([self.batch_size, self.seq_length]      )

     lowerCamelCase__       : Tuple   					=						GPTJConfig(
         vocab_size=self.vocab_size     ,	n_embd=self.hidden_size     ,	n_layer=self.num_hidden_layers     ,	n_head=self.num_attention_heads     ,	n_positions=self.max_position_embeddings     ,	use_cache=UpperCAmelCase     ,	bos_token_id=self.bos_token_id     ,	eos_token_id=self.eos_token_id     ,	pad_token_id=self.pad_token_id     ,	rotary_dim=self.rotary_dim     ,	)

     return (config, input_ids, input_mask)
 def 	A_      ( self    : Optional[Any]      )     ->      Union[str, Any]:
     lowerCamelCase__       : Tuple   					=						self.prepare_config_and_inputs()
     lowerCamelCase__					,				lowerCamelCase__					,				lowerCamelCase__       : List[Any]   					=						config_and_inputs
     lowerCamelCase__       : List[Any]   					=						{'input_ids': input_ids, 'attention_mask': attention_mask}
     return config, inputs_dict
 def 	A_      ( self    : int     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : Optional[Any]      )     ->      Tuple:
     lowerCamelCase__       : Tuple   					=						20
     lowerCamelCase__       : Dict   					=						model_class_name(UpperCAmelCase      )

     lowerCamelCase__       : Dict   					=						model.init_cache(input_ids.shape[0]     ,	UpperCAmelCase      )
     lowerCamelCase__       : int   					=						jnp.ones((input_ids.shape[0], max_decoder_length)     ,	dtype='i4'      )

     lowerCamelCase__       : List[str]   					=						jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1      )[None, :]     ,	(input_ids.shape[0], input_ids.shape[-1] - 1)      )
     lowerCamelCase__       : Optional[int]   					=						model(
         input_ids[:, :-1]     ,	attention_mask=UpperCAmelCase     ,	past_key_values=UpperCAmelCase     ,	position_ids=UpperCAmelCase     ,	)

     lowerCamelCase__       : List[str]   					=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]     ,	dtype='i4'      )
     lowerCamelCase__       : List[str]   					=						model(
         input_ids[:, -1:]     ,	attention_mask=UpperCAmelCase     ,	past_key_values=outputs_cache.past_key_values     ,	position_ids=UpperCAmelCase     ,	)

     lowerCamelCase__       : List[str]   					=						model(UpperCAmelCase      )

     lowerCamelCase__       : Optional[Any]   					=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
     self.parent.assertTrue(diff < 1e-3     ,	msg=F"""Max diff is {diff}"""      )







 def 	A_      ( self    : Union[str, Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : List[str]      )     ->      Optional[Any]:
     lowerCamelCase__       : Any   					=						20
     lowerCamelCase__       : Any   					=						model_class_name(UpperCAmelCase      )

     lowerCamelCase__       : List[str]   					=						jnp.concatenate(
         [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])      )]     ,	axis=-1     ,	)

     lowerCamelCase__       : str   					=						model.init_cache(input_ids.shape[0]     ,	UpperCAmelCase      )
     lowerCamelCase__       : Dict   					=						jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1      )[None, :]     ,	(input_ids.shape[0], input_ids.shape[-1] - 1)      )

     lowerCamelCase__       : List[Any]   					=						model(
         input_ids[:, :-1]     ,	attention_mask=UpperCAmelCase     ,	past_key_values=UpperCAmelCase     ,	position_ids=UpperCAmelCase     ,	)
     lowerCamelCase__       : Optional[Any]   					=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]     ,	dtype='i4'      )
     lowerCamelCase__       : List[Any]   					=						model(
         input_ids[:, -1:]     ,	past_key_values=outputs_cache.past_key_values     ,	attention_mask=UpperCAmelCase     ,	position_ids=UpperCAmelCase     ,	)

     lowerCamelCase__       : List[Any]   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase      )

     lowerCamelCase__       : Dict   					=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
     self.parent.assertTrue(diff < 1e-3     ,	msg=F"""Max diff is {diff}"""      )


@require_flax
class   lowerCAmelCase				(   __UpperCamelCase,    __UpperCamelCase,    unittest.TestCase     ):
 UpperCAmelCase__						=						(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
 UpperCAmelCase__						=						(FlaxGPTJForCausalLM,) if is_flax_available() else ()
 def 	A_      ( self    : Optional[Any]      )     ->      Tuple:
     lowerCamelCase__       : Optional[Any]   					=						FlaxGPTJModelTester(self      )
 def 	A_      ( self    : Tuple      )     ->      Union[str, Any]:
     for model_class_name in self.all_model_classes:
         lowerCamelCase__					,				lowerCamelCase__					,				lowerCamelCase__       : Optional[int]   					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward(UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase      )
 def 	A_      ( self    : Union[str, Any]      )     ->      Tuple:
     for model_class_name in self.all_model_classes:
         lowerCamelCase__					,				lowerCamelCase__					,				lowerCamelCase__       : str   					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward_with_attn_mask(
             UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase      )
 @tooslow
 def 	A_      ( self    : Dict      )     ->      Tuple:
     lowerCamelCase__       : str   					=						GPTaTokenizer.from_pretrained('gpt2'     ,	pad_token='<|endoftext|>'     ,	padding_side='left'      )
     lowerCamelCase__       : Dict   					=						tokenizer(['Hello this is a long string', 'Hey']     ,	return_tensors='np'     ,	padding=UpperCAmelCase     ,	truncation=UpperCAmelCase      )

     lowerCamelCase__       : str   					=						FlaxGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'      )
     lowerCamelCase__       : Optional[int]   					=						False
     lowerCamelCase__       : Optional[int]   					=						model.config.eos_token_id

     lowerCamelCase__       : Union[str, Any]   					=						jax.jit(model.generate      )

     lowerCamelCase__       : Optional[Any]   					=						jit_generate(
         inputs['input_ids']     ,	attention_mask=inputs['attention_mask']     ,	pad_token_id=tokenizer.pad_token_id      ).sequences

     lowerCamelCase__       : Optional[int]   					=						tokenizer.batch_decode(UpperCAmelCase     ,	skip_special_tokens=UpperCAmelCase      )

     lowerCamelCase__       : int   					=						[
         'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
         'Hey, I\'m a little late to the party. I\'m going to',
     ]

     self.assertListEqual(UpperCAmelCase     ,	UpperCAmelCase      )
 @is_pt_flax_cross_test
 def 	A_      ( self    : List[str]      )     ->      Union[str, Any]:
     lowerCamelCase__					,				lowerCamelCase__       : Optional[int]   					=						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__      ):
             # prepare inputs
             lowerCamelCase__       : Dict   					=						self._prepare_for_class(UpperCAmelCase     ,	UpperCAmelCase      )
             lowerCamelCase__       : Optional[int]   					=						{k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             lowerCamelCase__       : Any   					=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
             lowerCamelCase__       : Union[str, Any]   					=						getattr(UpperCAmelCase     ,	UpperCAmelCase      )

             lowerCamelCase__					,				lowerCamelCase__       : Tuple   					=						pt_inputs['input_ids'].shape
             lowerCamelCase__       : Tuple   					=						np.random.randint(0     ,	seq_length - 1     ,	size=(batch_size,)      )
             for batch_idx, start_index in enumerate(UpperCAmelCase      ):
                 lowerCamelCase__       : List[str]   					=						0
                 lowerCamelCase__       : Optional[int]   					=						1
                 lowerCamelCase__       : Any   					=						0
                 lowerCamelCase__       : Any   					=						1
             lowerCamelCase__       : Dict   					=						pt_model_class(UpperCAmelCase      ).eval()
             lowerCamelCase__       : str   					=						model_class(UpperCAmelCase     ,	dtype=jnp.floataa      )

             lowerCamelCase__       : str   					=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()     ,	UpperCAmelCase      )
             lowerCamelCase__       : Tuple   					=						fx_state

             with torch.no_grad():
                 lowerCamelCase__       : List[str]   					=						pt_model(**UpperCAmelCase      ).to_tuple()

             lowerCamelCase__       : Dict   					=						fx_model(**UpperCAmelCase      ).to_tuple()
             self.assertEqual(len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )

             with tempfile.TemporaryDirectory() as tmpdirname:
                 pt_model.save_pretrained(UpperCAmelCase      )
                 lowerCamelCase__       : Tuple   					=						model_class.from_pretrained(UpperCAmelCase     ,	from_pt=UpperCAmelCase      )

             lowerCamelCase__       : int   					=						fx_model_loaded(**UpperCAmelCase      ).to_tuple()
             self.assertEqual(
                 len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output_loaded, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output_loaded[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )
 @is_pt_flax_cross_test
 def 	A_      ( self    : Union[str, Any]      )     ->      int:
     lowerCamelCase__					,				lowerCamelCase__       : Dict   					=						self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__      ):
             # prepare inputs
             lowerCamelCase__       : Tuple   					=						self._prepare_for_class(UpperCAmelCase     ,	UpperCAmelCase      )
             lowerCamelCase__       : Tuple   					=						{k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             lowerCamelCase__       : Union[str, Any]   					=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
             lowerCamelCase__       : int   					=						getattr(UpperCAmelCase     ,	UpperCAmelCase      )

             lowerCamelCase__       : List[Any]   					=						pt_model_class(UpperCAmelCase      ).eval()
             lowerCamelCase__       : str   					=						model_class(UpperCAmelCase     ,	dtype=jnp.floataa      )

             lowerCamelCase__       : List[str]   					=						load_flax_weights_in_pytorch_model(UpperCAmelCase     ,	fx_model.params      )
             lowerCamelCase__					,				lowerCamelCase__       : Union[str, Any]   					=						pt_inputs['input_ids'].shape
             lowerCamelCase__       : Dict   					=						np.random.randint(0     ,	seq_length - 1     ,	size=(batch_size,)      )
             for batch_idx, start_index in enumerate(UpperCAmelCase      ):
                 lowerCamelCase__       : List[Any]   					=						0
                 lowerCamelCase__       : Tuple   					=						1
                 lowerCamelCase__       : Optional[Any]   					=						0
                 lowerCamelCase__       : List[str]   					=						1

             # make sure weights are tied in PyTorch
             pt_model.tie_weights()

             with torch.no_grad():
                 lowerCamelCase__       : Any   					=						pt_model(**UpperCAmelCase      ).to_tuple()

             lowerCamelCase__       : List[Any]   					=						fx_model(**UpperCAmelCase      ).to_tuple()
             self.assertEqual(len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )

             with tempfile.TemporaryDirectory() as tmpdirname:
                 fx_model.save_pretrained(UpperCAmelCase      )
                 lowerCamelCase__       : Tuple   					=						pt_model_class.from_pretrained(UpperCAmelCase     ,	from_flax=UpperCAmelCase      )

             with torch.no_grad():
                 lowerCamelCase__       : Optional[int]   					=						pt_model_loaded(**UpperCAmelCase      ).to_tuple()

             self.assertEqual(
                 len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )







 @tooslow
 def 	A_      ( self    : List[str]      )     ->      List[Any]:
     for model_class_name in self.all_model_classes:
         lowerCamelCase__       : Union[str, Any]   					=						model_class_name.from_pretrained('EleutherAI/gpt-j-6B'      )
         lowerCamelCase__       : Any   					=						model(np.ones((1, 1)      )      )
         self.assertIsNotNone(UpperCAmelCase      )





