






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"""







				UpperCamelCase__ :      int			 =XGLMConfig
				UpperCamelCase__ :      str			 ={}
				UpperCamelCase__ :      Dict			 ='''gelu'''


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




									"""simple docstring"""

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


				def       __lowercase							( self    ):




									"""simple docstring"""

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


				def       __lowercase							( self    ):




									"""simple docstring"""

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

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

									__UpperCamelCase      :					Dict       		=self.get_config()

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

									return (
									    config,
									    input_ids,
									    input_mask,
									    head_mask,
									)


				def       __lowercase							( self    ):




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







				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Tuple       		=self.prepare_config_and_inputs()

									(
									    (
									    __UpperCamelCase  
									)		,				(
									    __UpperCamelCase  
									)		,				(
									    __UpperCamelCase  
									)		,				(
									    __UpperCamelCase  
									)		,				  
									)      :					Tuple       		=config_and_inputs

									__UpperCamelCase      :					int       		={
									    'input_ids': input_ids,
									    'head_mask': head_mask,
									}

									return config, inputs_dict




@require_tf
class 				__A     (			__lowerCAmelCase      ,     __lowerCAmelCase      ,     unittest.TestCase ):





				"""simple docstring"""







				UpperCamelCase__ :      int			 =(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
				UpperCamelCase__ :      Union[str, Any]			 =(TFXGLMForCausalLM,) if is_tf_available() else ()
				UpperCamelCase__ :      List[str]			 =(
				    {'''feature-extraction''': TFXGLMModel, '''text-generation''': TFXGLMForCausalLM} if is_tf_available() else {}
				)
				UpperCamelCase__ :      Dict			 =False
				UpperCamelCase__ :      Optional[Any]			 =False
				UpperCamelCase__ :      Optional[int]			 =False


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					List[Any]       		=TFXGLMModelTester(self    )
									__UpperCamelCase      :					str       		=ConfigTester(self							,  config_class=lowerCAmelCase_							,  n_embd=37    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									self.config_tester.run_common_tests()


				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

									for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														__UpperCamelCase      :					Any       		=TFXGLMModel.from_pretrained(lowerCAmelCase_    )
														self.assertIsNotNone(lowerCAmelCase_    )







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




									"""simple docstring"""

									super().test_resize_token_embeddings()




@require_tf
class 				__A     (			unittest.TestCase ):





				"""simple docstring"""







				@slow
				def       __lowercase							( self							,  lowerCamelCase__=True    ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		=TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M'    )
									__UpperCamelCase      :					Optional[int]       		=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
									__UpperCamelCase      :					int       		=[2, 268, 9865, 67, 11, 1988, 57252, 9865, 5, 984, 67, 1988, 213838, 1658, 53, 70446, 33, 6657, 278, 1581]
									# fmt: on
									__UpperCamelCase      :					List[Any]       		=model.generate(lowerCAmelCase_							,  do_sample=lowerCAmelCase_							,  num_beams=1    )
									if verify_outputs:
														self.assertListEqual(output_ids[0].numpy().tolist()							,  lowerCAmelCase_    )


				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

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

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

									__UpperCamelCase      :					Optional[int]       		=(
									    'Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due'
									)
									self.assertEqual(lowerCAmelCase_							,  lowerCAmelCase_    )







				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

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

									__UpperCamelCase      :					Any       		='left'

									# use different length sentences to test batching
									__UpperCamelCase      :					Optional[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',
									]

									__UpperCamelCase      :					Optional[int]       		=tokenizer(lowerCAmelCase_							,  return_tensors='tf'							,  padding=lowerCAmelCase_    )
									__UpperCamelCase      :					List[str]       		=inputs['input_ids']

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

									__UpperCamelCase      :					str       		=tokenizer(sentences[0]							,  return_tensors='tf'    ).input_ids
									__UpperCamelCase      :					Optional[int]       		=model.generate(input_ids=lowerCAmelCase_							,  max_new_tokens=12    )

									__UpperCamelCase      :					Dict       		=tokenizer(sentences[1]							,  return_tensors='tf'    ).input_ids
									__UpperCamelCase      :					Optional[Any]       		=model.generate(input_ids=lowerCAmelCase_							,  max_new_tokens=12    )

									__UpperCamelCase      :					int       		=tokenizer.batch_decode(lowerCAmelCase_							,  skip_special_tokens=lowerCAmelCase_    )
									__UpperCamelCase      :					Dict       		=tokenizer.decode(output_non_padded[0]							,  skip_special_tokens=lowerCAmelCase_    )
									__UpperCamelCase      :					Any       		=tokenizer.decode(output_padded[0]							,  skip_special_tokens=lowerCAmelCase_    )

									__UpperCamelCase      :					Optional[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 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(lowerCAmelCase_							,  lowerCAmelCase_    )
									self.assertListEqual(lowerCAmelCase_							,  [non_padded_sentence, padded_sentence]    )


import torch
from torch import nn



class 				a							(			nn.Module       ):







 """simple docstring"""






 def __init__(					self		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_=1		,				lowerCAmelCase_=False							)  ->					Any:
        super().__init__()

        _A        =  n_token
        _A        =  d_embed
        _A        =  d_proj

        _A        =  cutoffs + [n_token]
        _A        =  [0] + self.cutoffs
        _A        =  div_val

        _A        =  self.cutoffs[0]
        _A        =  len(self.cutoffs							) - 1
        _A        =  self.shortlist_size + self.n_clusters

        if self.n_clusters > 0:
               _A        =  nn.Parameter(torch.zeros(self.n_clusters		,				self.d_embed							)							)
               _A        =  nn.Parameter(torch.zeros(self.n_clusters							)							)

        _A        =  nn.ModuleList()
        _A        =  nn.ParameterList()

        if div_val == 1:
               for i in range(len(self.cutoffs							)							):
                      if d_proj != d_embed:
                             self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCAmelCase_		,				lowerCAmelCase_							)							)							)
                      else:
                             self.out_projs.append(lowerCAmelCase_							)

               self.out_layers.append(nn.Linear(lowerCAmelCase_		,				lowerCAmelCase_							)							)
        else:
               for i in range(len(self.cutoffs							)							):
                      _A   ,  _A        =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                      _A        =  d_embed // (div_val**i)

                      self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCAmelCase_		,				lowerCAmelCase_							)							)							)

                      self.out_layers.append(nn.Linear(lowerCAmelCase_		,				r_idx - l_idx							)							)

        _A        =  keep_order



 def      UpperCAmelCase							(					self		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)  ->					Optional[Any]:
        if proj is None:
               _A        =  nn.functional.linear(lowerCAmelCase_		,				lowerCAmelCase_		,				bias=lowerCAmelCase_							)
        else:
               # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
               _A        =  nn.functional.linear(lowerCAmelCase_		,				proj.t().contiguous()							)
               _A        =  nn.functional.linear(lowerCAmelCase_		,				lowerCAmelCase_		,				bias=lowerCAmelCase_							)
               # else:
               #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
               #     if bias is not None:
               #         logit = logit + bias

        return logit



 def      UpperCAmelCase							(					self		,				lowerCAmelCase_		,				lowerCAmelCase_=None		,				lowerCAmelCase_=False							)  ->					List[Any]:

        if labels is not None:
               # Shift so that tokens < n predict n
               _A        =  hidden[..., :-1, :].contiguous()
               _A        =  labels[..., 1:].contiguous()
               _A        =  hidden.view(-1		,				hidden.size(-1							)							)
               _A        =  labels.view(-1							)
               if hidden.size(0							) != labels.size(0							):
                      raise RuntimeError("""Input and labels should have the same size in the batch dimension."""							)
        else:
               _A        =  hidden.view(-1		,				hidden.size(-1							)							)

        if self.n_clusters == 0:
               _A        =  self._compute_logit(lowerCAmelCase_		,				self.out_layers[0].weight		,				self.out_layers[0].bias		,				self.out_projs[0]							)
               if labels is not None:
                      _A        =  labels != -1_00
                      _A        =  torch.zeros_like(lowerCAmelCase_		,				dtype=hidden.dtype		,				device=hidden.device							)
                      _A        =  (
                          -nn.functional.log_softmax(lowerCAmelCase_		,				dim=-1							)[mask].gather(1		,				labels[mask].unsqueeze(1							)							).squeeze(1							)
                      )
               else:
                      _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=-1							)
        else:
               # construct weights and biases
               _A   ,  _A        =  [], []
               for i in range(len(self.cutoffs							)							):
                      if self.div_val == 1:
                             _A   ,  _A        =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             _A        =  self.out_layers[0].weight[l_idx:r_idx]
                             _A        =  self.out_layers[0].bias[l_idx:r_idx]
                      else:
                             _A        =  self.out_layers[i].weight
                             _A        =  self.out_layers[i].bias

                      if i == 0:
                             _A        =  torch.cat([weight_i, self.cluster_weight]		,				dim=0							)
                             _A        =  torch.cat([bias_i, self.cluster_bias]		,				dim=0							)

                      weights.append(lowerCAmelCase_							)
                      biases.append(lowerCAmelCase_							)

               _A   ,  _A   ,  _A        =  weights[0], biases[0], self.out_projs[0]

               _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)
               _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=1							)

               if labels is None:
                      _A        =  hidden.new_empty((head_logit.size(0							), self.n_token)							)
               else:
                      _A        =  torch.zeros_like(lowerCAmelCase_		,				dtype=hidden.dtype		,				device=hidden.device							)

               _A        =  0
               _A        =  [0] + self.cutoffs
               for i in range(len(lowerCAmelCase_							) - 1							):
                      _A   ,  _A        =  cutoff_values[i], cutoff_values[i + 1]

                      if labels is not None:
                             _A        =  (labels >= l_idx) & (labels < r_idx)
                             _A        =  mask_i.nonzero().squeeze()

                             if indices_i.numel() == 0:
                                    continue

                             _A        =  labels.index_select(0		,				lowerCAmelCase_							) - l_idx
                             _A        =  head_logprob.index_select(0		,				lowerCAmelCase_							)
                             _A        =  hidden.index_select(0		,				lowerCAmelCase_							)
                      else:
                             _A        =  hidden

                      if i == 0:
                             if labels is not None:
                                    _A        =  head_logprob_i.gather(1		,				target_i[:, None]							).squeeze(1							)
                             else:
                                    _A        =  head_logprob[:, : self.cutoffs[0]]
                      else:
                             _A   ,  _A   ,  _A        =  weights[i], biases[i], self.out_projs[i]

                             _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)
                             _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=1							)
                             _A        =  self.cutoffs[0] + i - 1  # No probability for the head cluster
                             if labels is not None:
                                    _A        =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                                        1		,				target_i[:, None]							).squeeze(1							)
                             else:
                                    _A        =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                                    _A        =  logprob_i

                      if labels is not None:
                             if (hasattr(self		,				"""keep_order"""							) and self.keep_order) or keep_order:
                                    out.index_copy_(0		,				lowerCAmelCase_		,				-logprob_i							)
                             else:
                                    out[offset : offset + logprob_i.size(0							)].copy_(-logprob_i							)
                             offset += logprob_i.size(0							)

        return out



 def      UpperCAmelCase							(					self		,				lowerCAmelCase_							)  ->					int:
        if self.n_clusters == 0:
               _A        =  self._compute_logit(lowerCAmelCase_		,				self.out_layers[0].weight		,				self.out_layers[0].bias		,				self.out_projs[0]							)
               return nn.functional.log_softmax(lowerCAmelCase_		,				dim=-1							)
        else:
               # construct weights and biases
               _A   ,  _A        =  [], []
               for i in range(len(self.cutoffs							)							):
                      if self.div_val == 1:
                             _A   ,  _A        =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             _A        =  self.out_layers[0].weight[l_idx:r_idx]
                             _A        =  self.out_layers[0].bias[l_idx:r_idx]
                      else:
                             _A        =  self.out_layers[i].weight
                             _A        =  self.out_layers[i].bias

                      if i == 0:
                             _A        =  torch.cat([weight_i, self.cluster_weight]		,				dim=0							)
                             _A        =  torch.cat([bias_i, self.cluster_bias]		,				dim=0							)

                      weights.append(lowerCAmelCase_							)
                      biases.append(lowerCAmelCase_							)

               _A   ,  _A   ,  _A        =  weights[0], biases[0], self.out_projs[0]
               _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)

               _A        =  hidden.new_empty((head_logit.size(0							), self.n_token)							)
               _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=1							)

               _A        =  [0] + self.cutoffs
               for i in range(len(lowerCAmelCase_							) - 1							):
                      _A   ,  _A        =  cutoff_values[i], cutoff_values[i + 1]

                      if i == 0:
                             _A        =  head_logprob[:, : self.cutoffs[0]]
                      else:
                             _A   ,  _A   ,  _A        =  weights[i], biases[i], self.out_projs[i]

                             _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)
                             _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=1							)

                             _A        =  head_logprob[:, -i] + tail_logprob_i
                             _A        =  logprob_i

               return out




