





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 					A__       :




   """simple docstring"""



   __A  :							Optional[Any]			    =    XGLMConfig
   __A  :							List[str]			    =    {}
   __A  :							Tuple			    =    '''gelu'''

   def __init__(   self   ,    lowercase   ,    lowercase=14   ,    lowercase=7   ,    lowercase=True   ,    lowercase=True   ,    lowercase=True   ,    lowercase=99   ,    lowercase=32   ,    lowercase=2   ,    lowercase=4   ,    lowercase=37   ,    lowercase="gelu"   ,    lowercase=0.1   ,    lowercase=0.1   ,    lowercase=512   ,    lowercase=0.02   ,    )     ->		int:







          '''simple docstring'''

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

   def 				__lowercase						(   self)     ->		Any:







          '''simple docstring'''

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

   def 				__lowercase						(   self)     ->		Any:







          '''simple docstring'''

          a__		:      List[str]					=							tf.clip_by_value(
              ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size)   ,    clip_value_min=0   ,    clip_value_max=3)

          a__		:      Dict					=							None
          if self.use_input_mask:
                 a__		:      Any					=							random_attention_mask([self.batch_size, self.seq_length])

          a__		:      Optional[int]					=							self.get_config()

          a__		:      Optional[int]					=							floats_tensor([self.num_hidden_layers, self.num_attention_heads]   ,    2)

          return (
              config,
              input_ids,
              input_mask,
              head_mask,
          )

   def 				__lowercase						(   self)     ->		Optional[Any]:







          '''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=__a   ,    bos_token_id=self.bos_token_id   ,    eos_token_id=self.eos_token_id   ,    pad_token_id=self.pad_token_id   ,    return_dict=__a   ,    )

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







          '''simple docstring'''

          a__		:      Any					=							self.prepare_config_and_inputs()

          (
              (
              a__				
          )					,				(
              a__				
          )					,				(
              a__				
          )					,				(
              a__				
          )					,								
          )		:      Tuple					=							config_and_inputs

          a__		:      Optional[int]					=							{
              'input_ids': input_ids,
              'head_mask': head_mask,
          }

          return config, inputs_dict





@require_tf
class 					A__       (				UpperCAmelCase_					,    UpperCAmelCase_					,    unittest.TestCase      ):




   """simple docstring"""



   __A  :							int			    =    (TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
   __A  :							int			    =    (TFXGLMForCausalLM,) if is_tf_available() else ()
   __A  :							Optional[int]			    =    (
       {'''feature-extraction''': TFXGLMModel, '''text-generation''': TFXGLMForCausalLM} if is_tf_available() else {}
   )
   __A  :							int			    =    False
   __A  :							Optional[int]			    =    False
   __A  :							Tuple			    =    False

   def 				__lowercase						(   self)     ->		Dict:







          '''simple docstring'''

          a__		:      str					=							TFXGLMModelTester(self)
          a__		:      Optional[Any]					=							ConfigTester(self   ,    config_class=__a   ,    n_embd=37)

   def 				__lowercase						(   self)     ->		Optional[Any]:







          '''simple docstring'''

          self.config_tester.run_common_tests()

   @slow
   def 				__lowercase						(   self)     ->		Tuple:







          '''simple docstring'''

          for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 a__		:      List[Any]					=							TFXGLMModel.from_pretrained(__a)
                 self.assertIsNotNone(__a)

   @unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.')
   def 				__lowercase						(   self)     ->		Optional[int]:







          '''simple docstring'''

          super().test_resize_token_embeddings()





@require_tf
class 					A__       (				unittest.TestCase      ):




   """simple docstring"""



   @slow
   def 				__lowercase						(   self   ,    lowercase=True)     ->		int:







          '''simple docstring'''

          a__		:      List[Any]					=							TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M')
          a__		:      List[str]					=							tf.convert_to_tensor([[2, 268, 9865]]   ,    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
          a__		:      Optional[Any]					=							[2, 268, 9865, 67, 11, 1988, 5_7252, 9865, 5, 984, 67, 1988, 21_3838, 1658, 53, 7_0446, 33, 6657, 278, 1581]
          # fmt: on
          a__		:      str					=							model.generate(__a   ,    do_sample=__a   ,    num_beams=1)
          if verify_outputs:
                 self.assertListEqual(output_ids[0].numpy().tolist()   ,    __a)

   @slow
   def 				__lowercase						(   self)     ->		str:







          '''simple docstring'''

          a__		:      Any					=							XGLMTokenizer.from_pretrained('facebook/xglm-564M')
          a__		:      Optional[Any]					=							TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M')

          tf.random.set_seed(0)
          a__		:      Any					=							tokenizer('Today is a nice day and'   ,    return_tensors='tf')
          a__		:      Union[str, Any]					=							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'):
                 a__		:      str					=							model.generate(__a   ,    do_sample=__a   ,    seed=[7, 0])
          a__		:      Optional[int]					=							tokenizer.decode(output_ids[0]   ,    skip_special_tokens=__a)

          a__		:      Tuple					=							(
              'Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due'
          )
          self.assertEqual(__a   ,    __a)

   @slow
   def 				__lowercase						(   self)     ->		int:







          '''simple docstring'''

          a__		:      Union[str, Any]					=							TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M')
          a__		:      Optional[Any]					=							XGLMTokenizer.from_pretrained('facebook/xglm-564M')

          a__		:      int					=							'left'

          # use different length sentences to test batching
          a__		:      List[str]					=							[
              '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',
          ]

          a__		:      Optional[Any]					=							tokenizer(__a   ,    return_tensors='tf'   ,    padding=__a)
          a__		:      int					=							inputs['input_ids']

          a__		:      Tuple					=							model.generate(input_ids=__a   ,    attention_mask=inputs['attention_mask']   ,    max_new_tokens=12)

          a__		:      Dict					=							tokenizer(sentences[0]   ,    return_tensors='tf').input_ids
          a__		:      List[str]					=							model.generate(input_ids=__a   ,    max_new_tokens=12)

          a__		:      Any					=							tokenizer(sentences[1]   ,    return_tensors='tf').input_ids
          a__		:      Optional[int]					=							model.generate(input_ids=__a   ,    max_new_tokens=12)

          a__		:      Tuple					=							tokenizer.batch_decode(__a   ,    skip_special_tokens=__a)
          a__		:      str					=							tokenizer.decode(output_non_padded[0]   ,    skip_special_tokens=__a)
          a__		:      Optional[Any]					=							tokenizer.decode(output_padded[0]   ,    skip_special_tokens=__a)

          a__		:      List[str]					=							[
              '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(__a   ,    __a)
          self.assertListEqual(__a   ,    [non_padded_sentence, padded_sentence])







import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest





class  lowerCamelCase__	(    UpperCAmelCase_):
     """simple docstring"""



     _A      =      (PNDMScheduler,)
     _A      =      (('num_inference_steps', 50),)







     def     _a   (self					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		{
               "num_train_timesteps": 10_00,
               "beta_start": 0.0001,
               "beta_end": 0.02,
               "beta_schedule": "linear",
           }

           config.update(**__a     )
           return config







     def     _a   (self					,      __a=0					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		dict(self.forward_default_kwargs     )
           lowerCamelCase             =		kwargs.pop("num_inference_steps"					,      __a     )
           lowerCamelCase             =		self.dummy_sample
           lowerCamelCase             =		0.1 * sample
           lowerCamelCase             =		[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.get_scheduler_config(**__a     )
                 lowerCamelCase             =		scheduler_class(**__a     )
                 scheduler.set_timesteps(__a     )
                 # copy over dummy past residuals
                 lowerCamelCase             =		dummy_past_residuals[:]

                 with tempfile.TemporaryDirectory() as tmpdirname:
                       scheduler.save_config(__a     )
                       lowerCamelCase             =		scheduler_class.from_pretrained(__a     )
                       new_scheduler.set_timesteps(__a     )
                       # copy over dummy past residuals
                       lowerCamelCase             =		dummy_past_residuals[:]

                 lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"

                 lowerCamelCase             =		scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"







     def     _a   (self     ):






           '''simple docstring'''



           pass







     def     _a   (self					,      __a=0					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		dict(self.forward_default_kwargs     )
           lowerCamelCase             =		kwargs.pop("num_inference_steps"					,      __a     )
           lowerCamelCase             =		self.dummy_sample
           lowerCamelCase             =		0.1 * sample
           lowerCamelCase             =		[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )
                 scheduler.set_timesteps(__a     )

                 # copy over dummy past residuals (must be after setting timesteps)
                 lowerCamelCase             =		dummy_past_residuals[:]

                 with tempfile.TemporaryDirectory() as tmpdirname:
                       scheduler.save_config(__a     )
                       lowerCamelCase             =		scheduler_class.from_pretrained(__a     )
                       # copy over dummy past residuals
                       new_scheduler.set_timesteps(__a     )

                       # copy over dummy past residual (must be after setting timesteps)
                       lowerCamelCase             =		dummy_past_residuals[:]

                 lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"

                 lowerCamelCase             =		scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"







     def     _a   (self					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		self.scheduler_classes[0]
           lowerCamelCase             =		self.get_scheduler_config(**__a     )
           lowerCamelCase             =		scheduler_class(**__a     )

           lowerCamelCase             =		10
           lowerCamelCase             =		self.dummy_model()
           lowerCamelCase             =		self.dummy_sample_deter
           scheduler.set_timesteps(__a     )

           for i, t in enumerate(scheduler.prk_timesteps     ):
                 lowerCamelCase             =		model(__a					,      __a     )
                 lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a     ).prev_sample

           for i, t in enumerate(scheduler.plms_timesteps     ):
                 lowerCamelCase             =		model(__a					,      __a     )
                 lowerCamelCase             =		scheduler.step_plms(__a					,      __a					,      __a     ).prev_sample

           return sample







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		dict(self.forward_default_kwargs     )

           lowerCamelCase             =		kwargs.pop("num_inference_steps"					,      __a     )

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )

                 lowerCamelCase             =		self.dummy_sample
                 lowerCamelCase             =		0.1 * sample

                 if num_inference_steps is not None and hasattr(__a					,      "set_timesteps"     ):
                       scheduler.set_timesteps(__a     )
                 elif num_inference_steps is not None and not hasattr(__a					,      "set_timesteps"     ):
                       lowerCamelCase             =		num_inference_steps

                 # copy over dummy past residuals (must be done after set_timesteps)
                 lowerCamelCase             =		[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
                 lowerCamelCase             =		dummy_past_residuals[:]

                 lowerCamelCase             =		scheduler.step_prk(__a					,      0					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		scheduler.step_prk(__a					,      1					,      __a					,      **__a     ).prev_sample

                 self.assertEqual(output_a.shape					,      sample.shape     )
                 self.assertEqual(output_a.shape					,      output_a.shape     )

                 lowerCamelCase             =		scheduler.step_plms(__a					,      0					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		scheduler.step_plms(__a					,      1					,      __a					,      **__a     ).prev_sample

                 self.assertEqual(output_a.shape					,      sample.shape     )
                 self.assertEqual(output_a.shape					,      output_a.shape     )







     def     _a   (self     ):






           '''simple docstring'''



           for timesteps in [1_00, 10_00]:
                 self.check_over_configs(num_train_timesteps=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for steps_offset in [0, 1]:
                 self.check_over_configs(steps_offset=__a     )

           lowerCamelCase             =		self.scheduler_classes[0]
           lowerCamelCase             =		self.get_scheduler_config(steps_offset=1     )
           lowerCamelCase             =		scheduler_class(**__a     )
           scheduler.set_timesteps(10     )
           assert torch.equal(
               scheduler.timesteps					,      torch.LongTensor(
                   [9_01, 8_51, 8_51, 8_01, 8_01, 7_51, 7_51, 7_01, 7_01, 6_51, 6_51, 6_01, 6_01, 5_01, 4_01, 3_01, 2_01, 1_01, 1]     )					,      )







     def     _a   (self     ):






           '''simple docstring'''



           for beta_start, beta_end in zip([0.0001, 0.001]					,      [0.002, 0.02]     ):
                 self.check_over_configs(beta_start=__a					,      beta_end=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for schedule in ["linear", "squaredcos_cap_v2"]:
                 self.check_over_configs(beta_schedule=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for prediction_type in ["epsilon", "v_prediction"]:
                 self.check_over_configs(prediction_type=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for t in [1, 5, 10]:
                 self.check_over_forward(time_step=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for t, num_inference_steps in zip([1, 5, 10]					,      [10, 50, 1_00]     ):
                 self.check_over_forward(num_inference_steps=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		27

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.dummy_sample
                 lowerCamelCase             =		0.1 * sample

                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )

                 scheduler.set_timesteps(__a     )

                 # before power of 3 fix, would error on first step, so we only need to do two
                 for i, t in enumerate(scheduler.prk_timesteps[:2]     ):
                       lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a     ).prev_sample







     def     _a   (self     ):






           '''simple docstring'''



           with self.assertRaises(__a     ):
                 lowerCamelCase             =		self.scheduler_classes[0]
                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )

                 scheduler.step_plms(self.dummy_sample					,      1					,      self.dummy_sample     ).prev_sample







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop()
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 198.1318     ) < 1E-2
           assert abs(result_mean.item() - 0.2580     ) < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop(prediction_type="v_prediction"     )
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 67.3986     ) < 1E-2
           assert abs(result_mean.item() - 0.0878     ) < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop(set_alpha_to_one=__a					,      beta_start=0.01     )
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 230.0399     ) < 1E-2
           assert abs(result_mean.item() - 0.2995     ) < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop(set_alpha_to_one=__a					,      beta_start=0.01     )
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 186.9482     ) < 1E-2
           assert abs(result_mean.item() - 0.2434     ) < 1E-3