





from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers.models.xglm.modeling_tf_xglm import (
         TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFXGLMForCausalLM,
         TFXGLMModel,
     )



@require_tf
class    _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''
 lowercase_		=			XGLMConfig
 lowercase_		=			{}
 lowercase_		=			"gelu"

 def __init__(self				:   Optional[Any]       , UpperCAmelCase_				:   Union[str, Any]       , UpperCAmelCase_				:   Optional[int]=14       , UpperCAmelCase_				:   int=7       , UpperCAmelCase_				:   str=True       , UpperCAmelCase_				:   Optional[int]=True       , UpperCAmelCase_				:   int=True       , UpperCAmelCase_				:   Dict=99       , UpperCAmelCase_				:   Optional[int]=32       , UpperCAmelCase_				:   str=2       , UpperCAmelCase_				:   List[str]=4       , UpperCAmelCase_				:   List[str]=37       , UpperCAmelCase_				:   int="gelu"       , UpperCAmelCase_				:   Tuple=0.1       , UpperCAmelCase_				:   Union[str, Any]=0.1       , UpperCAmelCase_				:   Any=512       , UpperCAmelCase_				:   List[str]=0.02       , )       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =parent
   lowerCamelCase__:      int			  =batch_size
   lowerCamelCase__:      List[str]			  =seq_length
   lowerCamelCase__:      int			  =is_training
   lowerCamelCase__:      int			  =use_input_mask
   lowerCamelCase__:      Tuple			  =use_labels
   lowerCamelCase__:      Any			  =vocab_size
   lowerCamelCase__:      Dict			  =d_model
   lowerCamelCase__:      Optional[int]			  =num_hidden_layers
   lowerCamelCase__:      Tuple			  =num_attention_heads
   lowerCamelCase__:      Dict			  =ffn_dim
   lowerCamelCase__:      Optional[int]			  =activation_function
   lowerCamelCase__:      str			  =activation_dropout
   lowerCamelCase__:      int			  =attention_dropout
   lowerCamelCase__:      Optional[Any]			  =max_position_embeddings
   lowerCamelCase__:      Optional[int]			  =initializer_range
   lowerCamelCase__:      str			  =None
   lowerCamelCase__:      List[Any]			  =0
   lowerCamelCase__:      str			  =2
   lowerCamelCase__:      Union[str, Any]			  =1

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->Union[str, Any]:



   '''simple docstring'''



   return XGLMConfig.from_pretrained("facebook/xglm-564M")

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->int:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =tf.clip_by_value(
       ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size)       , clip_value_min=0       , clip_value_max=3)

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

   lowerCamelCase__:      Dict			  =self.get_config()

   lowerCamelCase__:      Optional[Any]			  =floats_tensor([self.num_hidden_layers, self.num_attention_heads]       , 2)

   return (
       config,
       input_ids,
       input_mask,
       head_mask,
   )

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->int:



   '''simple docstring'''



   return XGLMConfig(
       vocab_size=self.vocab_size       , d_model=self.hidden_size       , num_layers=self.num_hidden_layers       , attention_heads=self.num_attention_heads       , ffn_dim=self.ffn_dim       , activation_function=self.activation_function       , activation_dropout=self.activation_dropout       , attention_dropout=self.attention_dropout       , max_position_embeddings=self.max_position_embeddings       , initializer_range=self.initializer_range       , use_cache=UpperCAmelCase_       , bos_token_id=self.bos_token_id       , eos_token_id=self.eos_token_id       , pad_token_id=self.pad_token_id       , return_dict=UpperCAmelCase_       , )







 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      str			  =self.prepare_config_and_inputs()

   (
       (
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					
   ):      Union[str, Any]			  =config_and_inputs

   lowerCamelCase__:      str			  ={
       "input_ids": input_ids,
       "head_mask": head_mask,
   }

   return config, inputs_dict



@require_tf
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
 lowercase_		=			(TFXGLMForCausalLM,) if is_tf_available() else ()
 lowercase_		=			(
     {"feature-extraction": TFXGLMModel, "text-generation": TFXGLMForCausalLM} if is_tf_available() else {}
 )
 lowercase_		=			False
 lowercase_		=			False
 lowercase_		=			False

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =TFXGLMModelTester(self)
   lowerCamelCase__:      str			  =ConfigTester(self       , config_class=UpperCAmelCase_       , n_embd=37)

 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->str:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->Dict:



   '''simple docstring'''



   for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__:      Dict			  =TFXGLMModel.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)







 @unittest.skip(reason="Currently, model embeddings are going to undergo a major refactor.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Dict:



   '''simple docstring'''



   super().test_resize_token_embeddings()



@require_tf
class    _SCREAMING_SNAKE_CASE						(	unittest.TestCase    ):







 '''simple docstring'''

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any       , UpperCAmelCase_				:   str=True)       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =TFXGLMForCausalLM.from_pretrained("facebook/xglm-564M")
   lowerCamelCase__:      Optional[int]			  =tf.convert_to_tensor([[2, 268, 9_865]]       , dtype=tf.intaa)  # The dog
   # </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
   # fmt: off
   lowerCamelCase__:      Union[str, Any]			  =[2, 268, 9_865, 67, 11, 1_988, 57_252, 9_865, 5, 984, 67, 1_988, 213_838, 1_658, 53, 70_446, 33, 6_657, 278, 1_581]
   # fmt: on
   lowerCamelCase__:      List[str]			  =model.generate(UpperCAmelCase_       , do_sample=UpperCAmelCase_       , num_beams=1)
   if verify_outputs:
     self.assertListEqual(output_ids[0].numpy().tolist()       , UpperCAmelCase_)

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =XGLMTokenizer.from_pretrained("facebook/xglm-564M")
   lowerCamelCase__:      Tuple			  =TFXGLMForCausalLM.from_pretrained("facebook/xglm-564M")

   tf.random.set_seed(0)
   lowerCamelCase__:      Union[str, Any]			  =tokenizer("Today is a nice day and"       , return_tensors="tf")
   lowerCamelCase__:      Dict			  =tokenized.input_ids
   # forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
   with tf.device(":/CPU:0"):
     lowerCamelCase__:      str			  =model.generate(UpperCAmelCase_       , do_sample=UpperCAmelCase_       , seed=[7, 0])
   lowerCamelCase__:      str			  =tokenizer.decode(output_ids[0]       , skip_special_tokens=UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =(
       "Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due"
   )
   self.assertEqual(UpperCAmelCase_       , UpperCAmelCase_)







 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =TFXGLMForCausalLM.from_pretrained("facebook/xglm-564M")
   lowerCamelCase__:      Tuple			  =XGLMTokenizer.from_pretrained("facebook/xglm-564M")

   lowerCamelCase__:      str			  ="left"

   # use different length sentences to test batching
   lowerCamelCase__:      List[Any]			  =[
       "This is an extremelly long sentence that only exists to test the ability of the model to cope with "
       "left-padding, such as in batched generation. The output for the sequence below should be the same "
       "regardless of whether left padding is applied or not. When",
       "Hello, my dog is a little",
   ]

   lowerCamelCase__:      Optional[Any]			  =tokenizer(UpperCAmelCase_       , return_tensors="tf"       , padding=UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =inputs["input_ids"]

   lowerCamelCase__:      str			  =model.generate(input_ids=UpperCAmelCase_       , attention_mask=inputs["attention_mask"]       , max_new_tokens=12)

   lowerCamelCase__:      Tuple			  =tokenizer(sentences[0]       , return_tensors="tf").input_ids
   lowerCamelCase__:      Union[str, Any]			  =model.generate(input_ids=UpperCAmelCase_       , max_new_tokens=12)

   lowerCamelCase__:      Optional[Any]			  =tokenizer(sentences[1]       , return_tensors="tf").input_ids
   lowerCamelCase__:      Union[str, Any]			  =model.generate(input_ids=UpperCAmelCase_       , max_new_tokens=12)

   lowerCamelCase__:      Optional[Any]			  =tokenizer.batch_decode(UpperCAmelCase_       , skip_special_tokens=UpperCAmelCase_)
   lowerCamelCase__:      Any			  =tokenizer.decode(output_non_padded[0]       , skip_special_tokens=UpperCAmelCase_)
   lowerCamelCase__:      Union[str, Any]			  =tokenizer.decode(output_padded[0]       , skip_special_tokens=UpperCAmelCase_)

   lowerCamelCase__:      int			  =[
       "This is an extremelly long sentence that only exists to test the ability of the model to cope with "
       "left-padding, such as in batched generation. The output for the sequence below should be the same "
       "regardless of whether left padding is applied or not. When left padding is applied, the sequence will be "
       "a single",
       "Hello, my dog is a little bit of a shy one, but he is very friendly",
   ]
   self.assertListEqual(UpperCAmelCase_       , UpperCAmelCase_)
   self.assertListEqual(UpperCAmelCase_       , [non_padded_sentence, padded_sentence])


'''simple docstring'''




# Copyright 2021 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 warnings
from typing import List
from unittest.mock import Mock

import torch
from torch.utils.data import DataLoader, IterableDataset, TensorDataset

from accelerate.accelerator import Accelerator
from accelerate.utils.dataclasses import DistributedType





class     __magic_name__						(    _UpperCAmelCase):







			def __init__(	self				:  Optional[Any]				,			lowercase_				:  str    ):
										lowercase_					:			int										=       data







			def __iter__(	self				:  int    ):
										for element in self.data:
																	yield element




def 							lowerCamelCase	(      UpperCAmelCase__       :					Any=True )			->      Any:
							lowercase_					:			Optional[int]										=       Accelerator(even_batches=UpperCAmelCase__ )
							assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
							return accelerator




def 							lowerCamelCase	(      UpperCAmelCase__       :					Accelerator      ,      UpperCAmelCase__       :					int      ,      UpperCAmelCase__       :					int      ,      UpperCAmelCase__       :					bool = False )			->      Optional[Any]:
							if iterable:
														lowercase_					:			Dict										=       DummyIterableDataset(torch.as_tensor(range(UpperCAmelCase__ ) ) )
							else:
														lowercase_					:			Union[str, Any]										=       TensorDataset(torch.as_tensor(range(UpperCAmelCase__ ) ) )

							lowercase_					:			Any										=       DataLoader(UpperCAmelCase__      ,      batch_size=UpperCAmelCase__ )
							lowercase_					:			Any										=       accelerator.prepare(UpperCAmelCase__ )

							return dl




def 							lowerCamelCase	(      UpperCAmelCase__       :					Accelerator      ,      UpperCAmelCase__       :					int      ,      UpperCAmelCase__       :					int      ,      UpperCAmelCase__       :					List[int]      ,      UpperCAmelCase__       :					List[int]      ,      )			->      int:
							lowercase_					:			List[str]										=       create_dataloader(accelerator=UpperCAmelCase__      ,      dataset_size=UpperCAmelCase__      ,      batch_size=UpperCAmelCase__ )

							lowercase_					:			Tuple										=       [len(batch[0] ) for batch in dl]

							if accelerator.process_index == 0:
														assert batch_sizes == process_0_expected_batch_sizes
							elif accelerator.process_index == 1:
														assert batch_sizes == process_1_expected_batch_sizes




def 							lowerCamelCase	(      )			->      int:
							lowercase_					:			Tuple										=       create_accelerator()

							# without padding, we would expect a different number of batches
							verify_dataloader_batch_sizes(
							    UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1      ,      process_0_expected_batch_sizes=[1, 1]      ,      process_1_expected_batch_sizes=[1, 1]      ,      )

							# without padding, we would expect the same number of batches, but different sizes
							verify_dataloader_batch_sizes(
							    UpperCAmelCase__      ,      dataset_size=7      ,      batch_size=2      ,      process_0_expected_batch_sizes=[2, 2]      ,      process_1_expected_batch_sizes=[2, 2]      ,      )




def 							lowerCamelCase	(      )			->      Optional[int]:
							lowercase_					:			Optional[int]										=       create_accelerator(even_batches=UpperCAmelCase__ )

							verify_dataloader_batch_sizes(
							    UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1      ,      process_0_expected_batch_sizes=[1, 1]      ,      process_1_expected_batch_sizes=[1]      ,      )

							verify_dataloader_batch_sizes(
							    UpperCAmelCase__      ,      dataset_size=7      ,      batch_size=2      ,      process_0_expected_batch_sizes=[2, 2]      ,      process_1_expected_batch_sizes=[2, 1]      ,      )




def 							lowerCamelCase	(      )			->      List[str]:
							lowercase_					:			str										=       create_accelerator(even_batches=UpperCAmelCase__ )

							lowercase_					:			Dict										=       torch.nn.Linear(1      ,      1 )
							lowercase_					:			Optional[Any]										=       accelerator.prepare(UpperCAmelCase__ )

							lowercase_					:			List[Any]										=       create_dataloader(UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1 )

							lowercase_					:			Optional[Any]										=       []
							with accelerator.join_uneven_inputs([ddp_model] ):
														for batch_idx, batch in enumerate(UpperCAmelCase__ ):
																					lowercase_					:			Any										=       ddp_model(batch[0].float() )
																					lowercase_					:			List[str]										=       output.sum()
																					loss.backward()
																					batch_idxs.append(UpperCAmelCase__ )

							accelerator.wait_for_everyone()

							if accelerator.process_index == 0:
														assert batch_idxs == [0, 1]
							elif accelerator.process_index == 1:
														assert batch_idxs == [0]




def 							lowerCamelCase	(      UpperCAmelCase__       :					List[str] )			->      List[str]:
							with warnings.catch_warnings(record=UpperCAmelCase__ ) as w:
														with accelerator.join_uneven_inputs([Mock()] ):
																					pass

														assert issubclass(w[-1].category      ,      UpperCAmelCase__ )
														assert "only supported for multi-GPU" in str(w[-1].message )




def 							lowerCamelCase	(      )			->      Any:
							lowercase_					:			str										=       True
							lowercase_					:			Tuple										=       False
							lowercase_					:			str										=       create_accelerator(even_batches=UpperCAmelCase__ )
							lowercase_					:			Union[str, Any]										=       torch.nn.Linear(1      ,      1 )
							lowercase_					:			Any										=       accelerator.prepare(UpperCAmelCase__ )
							lowercase_					:			Optional[int]										=       create_dataloader(UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1 )
							lowercase_					:			List[Any]										=       create_dataloader(UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1 )

							with accelerator.join_uneven_inputs([ddp_model]      ,      even_batches=UpperCAmelCase__ ):
														lowercase_					:			Union[str, Any]										=       train_dl.batch_sampler.even_batches
														lowercase_					:			List[str]										=       valid_dl.batch_sampler.even_batches

							assert train_dl_overridden_value == overridden_even_batches
							assert valid_dl_overridden_value == overridden_even_batches
							assert train_dl.batch_sampler.even_batches == default_even_batches
							assert valid_dl.batch_sampler.even_batches == default_even_batches




def 							lowerCamelCase	(      )			->      Dict:
							lowercase_					:			str										=       True
							lowercase_					:			Optional[Any]										=       False
							lowercase_					:			Union[str, Any]										=       create_accelerator(even_batches=UpperCAmelCase__ )
							lowercase_					:			Optional[int]										=       torch.nn.Linear(1      ,      1 )
							lowercase_					:			Optional[int]										=       accelerator.prepare(UpperCAmelCase__ )
							create_dataloader(UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1      ,      iterable=UpperCAmelCase__ )
							lowercase_					:			List[str]										=       create_dataloader(UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1 )

							with warnings.catch_warnings():
														warnings.filterwarnings("""ignore""" )
														try:
																					with accelerator.join_uneven_inputs([ddp_model]      ,      even_batches=UpperCAmelCase__ ):
																												lowercase_					:			Optional[Any]										=       batch_dl.batch_sampler.even_batches
														except AttributeError:
																					# ensure attribute error is not raised when processing iterable dl
																					raise AssertionError

							assert batch_dl_overridden_value == overridden_even_batches
							assert batch_dl.batch_sampler.even_batches == default_even_batches




def 							lowerCamelCase	(      )			->      List[Any]:
							lowercase_					:			Optional[Any]										=       create_accelerator()
							lowercase_					:			Optional[int]										=       torch.nn.Linear(1      ,      1 )
							lowercase_					:			List[Any]										=       accelerator.prepare(UpperCAmelCase__ )
							create_dataloader(UpperCAmelCase__      ,      dataset_size=3      ,      batch_size=1      ,      iterable=UpperCAmelCase__ )

							with warnings.catch_warnings(record=UpperCAmelCase__ ) as w:
														with accelerator.join_uneven_inputs([ddp_model]      ,      even_batches=UpperCAmelCase__ ):
																					pass

														assert issubclass(w[-1].category      ,      UpperCAmelCase__ )
														assert "only supported for map-style datasets" in str(w[-1].message )




def 							lowerCamelCase	(      )			->      List[str]:
							lowercase_					:			List[Any]										=       create_accelerator()

							accelerator.print("""Test that even_batches variable ensures uniform batches across processes""" )
							test_default_ensures_even_batch_sizes()

							accelerator.print("""Run tests with even_batches disabled""" )
							test_can_disable_even_batches()

							accelerator.print("""Test joining uneven inputs""" )
							test_can_join_uneven_inputs()

							accelerator.print("""Test overriding even_batches when joining uneven inputs""" )
							test_join_can_override_even_batches()

							accelerator.print("""Test overriding even_batches for mixed dataloader types""" )
							test_join_can_override_for_mixed_type_dataloaders()

							accelerator.print("""Test overriding even_batches raises a warning for iterable dataloaders""" )
							test_join_raises_warning_for_iterable_when_overriding_even_batches()

							accelerator.print("""Test join with non DDP distributed raises warning""" )
							lowercase_					:			List[Any]										=       accelerator.state.distributed_type
							lowercase_					:			Union[str, Any]										=       DistributedType.FSDP
							test_join_raises_warning_for_non_ddp_distributed(UpperCAmelCase__ )
							lowercase_					:			str										=       original_state


if __name__ == "__main__":
						main()





