




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						       =						XGLMConfig
 UpperCamelCase						       =						{}
 UpperCamelCase						       =						'''gelu'''

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

       """simple docstring"""






       lowerCamelCase_							=   parent
       lowerCamelCase_							=   batch_size
       lowerCamelCase_							=   seq_length
       lowerCamelCase_							=   is_training
       lowerCamelCase_							=   use_input_mask
       lowerCamelCase_							=   use_labels
       lowerCamelCase_							=   vocab_size
       lowerCamelCase_							=   d_model
       lowerCamelCase_							=   num_hidden_layers
       lowerCamelCase_							=   num_attention_heads
       lowerCamelCase_							=   ffn_dim
       lowerCamelCase_							=   activation_function
       lowerCamelCase_							=   activation_dropout
       lowerCamelCase_							=   attention_dropout
       lowerCamelCase_							=   max_position_embeddings
       lowerCamelCase_							=   initializer_range
       lowerCamelCase_							=   None
       lowerCamelCase_							=   0
       lowerCamelCase_							=   2
       lowerCamelCase_							=   1

 def     a__			(    self :  List[Any]  )       ->     int:

       """simple docstring"""






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

 def     a__			(    self :  Dict  )       ->     List[str]:

       """simple docstring"""






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

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

       lowerCamelCase_							=   self.get_config()

       lowerCamelCase_							=   floats_tensor([self.num_hidden_layers, self.num_attention_heads]						,							2  )

       return (
           config,
           input_ids,
           input_mask,
           head_mask,
       )

 def     a__			(    self :  int  )       ->     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=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     a__			(    self :  List[str]  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   self.prepare_config_and_inputs()

       (
           (
           lowerCamelCase_
       )					,    (
           lowerCamelCase_
       )					,    (
           lowerCamelCase_
       )					,    (
           lowerCamelCase_
       )					,    
       )							=   config_and_inputs

       lowerCamelCase_							=   {
           'input_ids': input_ids,
           'head_mask': head_mask,
       }

       return config, inputs_dict







@require_tf
class   A(						UpperCamelCase							,				UpperCamelCase							,				unittest.TestCase       ):



 '''simple docstring'''


 UpperCamelCase						       =						(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
 UpperCamelCase						       =						(TFXGLMForCausalLM,) if is_tf_available() else ()
 UpperCamelCase						       =						(
     {'''feature-extraction''': TFXGLMModel, '''text-generation''': TFXGLMForCausalLM} if is_tf_available() else {}
 )
 UpperCamelCase						       =						False
 UpperCamelCase						       =						False
 UpperCamelCase						       =						False

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

       """simple docstring"""






       lowerCamelCase_							=   TFXGLMModelTester(self  )
       lowerCamelCase_							=   ConfigTester(self						,							config_class=A_						,							n_embd=37  )

 def     a__			(    self :  Optional[Any]  )       ->     Dict:

       """simple docstring"""






       self.config_tester.run_common_tests()

 @slow
 def     a__			(    self :  Optional[Any]  )       ->     Union[str, Any]:

       """simple docstring"""






       for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCamelCase_							=   TFXGLMModel.from_pretrained(A_  )
             self.assertIsNotNone(A_  )





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

       """simple docstring"""






       super().test_resize_token_embeddings()







@require_tf
class   A(						unittest.TestCase       ):



 '''simple docstring'''


 @slow
 def     a__			(    self :  int						,							A_ :  Dict=True  )       ->     Any:

       """simple docstring"""






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

 @slow
 def     a__			(    self :  Optional[Any]  )       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   XGLMTokenizer.from_pretrained('facebook/xglm-564M'  )
       lowerCamelCase_							=   TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M'  )

       tf.random.set_seed(0  )
       lowerCamelCase_							=   tokenizer('Today is a nice day and'						,							return_tensors='tf'  )
       lowerCamelCase_							=   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_							=   model.generate(A_						,							do_sample=A_						,							seed=[7, 0]  )
       lowerCamelCase_							=   tokenizer.decode(output_ids[0]						,							skip_special_tokens=A_  )

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





 @slow
 def     a__			(    self :  Dict  )       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M'  )
       lowerCamelCase_							=   XGLMTokenizer.from_pretrained('facebook/xglm-564M'  )

       lowerCamelCase_							=   'left'

       # use different length sentences to test batching
       lowerCamelCase_							=   [
           '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_							=   tokenizer(A_						,							return_tensors='tf'						,							padding=A_  )
       lowerCamelCase_							=   inputs['input_ids']

       lowerCamelCase_							=   model.generate(input_ids=A_						,							attention_mask=inputs['attention_mask']						,							max_new_tokens=12  )

       lowerCamelCase_							=   tokenizer(sentences[0]						,							return_tensors='tf'  ).input_ids
       lowerCamelCase_							=   model.generate(input_ids=A_						,							max_new_tokens=12  )

       lowerCamelCase_							=   tokenizer(sentences[1]						,							return_tensors='tf'  ).input_ids
       lowerCamelCase_							=   model.generate(input_ids=A_						,							max_new_tokens=12  )

       lowerCamelCase_							=   tokenizer.batch_decode(A_						,							skip_special_tokens=A_  )
       lowerCamelCase_							=   tokenizer.decode(output_non_padded[0]						,							skip_special_tokens=A_  )
       lowerCamelCase_							=   tokenizer.decode(output_padded[0]						,							skip_special_tokens=A_  )

       lowerCamelCase_							=   [
           '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]  )









from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin







@dataclass
class   A(						UpperCamelCase       ):



 '''simple docstring'''


 UpperCamelCase						       =						42







class   A(						UpperCamelCase							,				UpperCamelCase       ):



 '''simple docstring'''


 @register_to_config
 def __init__(    self :  Tuple						,							A_ :  int = 32						,							A_ :  int = 64						,							A_ :  int = 20						,							A_ :  int = 768						,							A_ :  Optional[Any]=77						,							A_ :  Optional[int]=4						,							A_ :  float = 0.0						,							A_ :  str = "silu"						,							A_ :  Optional[str] = None						,							A_ :  Optional[str] = None						,							A_ :  Optional[str] = "linear"						,							A_ :  Optional[str] = "prd"						,							A_ :  Optional[int] = None						,							A_ :  Optional[int] = None						,							A_ :  Optional[int] = None						,							)       ->     List[Any]:

       """simple docstring"""






       super().__init__()
       lowerCamelCase_							=   num_attention_heads
       lowerCamelCase_							=   attention_head_dim
       lowerCamelCase_							=   num_attention_heads * attention_head_dim
       lowerCamelCase_							=   additional_embeddings

       lowerCamelCase_							=   time_embed_dim or inner_dim
       lowerCamelCase_							=   embedding_proj_dim or embedding_dim
       lowerCamelCase_							=   clip_embed_dim or embedding_dim

       lowerCamelCase_							=   Timesteps(A_						,							A_						,							0  )
       lowerCamelCase_							=   TimestepEmbedding(A_						,							A_						,							out_dim=A_						,							act_fn=A_  )

       lowerCamelCase_							=   nn.Linear(A_						,							A_  )

       if embedding_proj_norm_type is None:
             lowerCamelCase_							=   None
       elif embedding_proj_norm_type == "layer":
             lowerCamelCase_							=   nn.LayerNorm(A_  )
       else:
             raise ValueError(f"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""  )

       lowerCamelCase_							=   nn.Linear(A_						,							A_  )

       if encoder_hid_proj_type is None:
             lowerCamelCase_							=   None
       elif encoder_hid_proj_type == "linear":
             lowerCamelCase_							=   nn.Linear(A_						,							A_  )
       else:
             raise ValueError(f"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""  )

       lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							num_embeddings + additional_embeddings						,							A_  )  )

       if added_emb_type == "prd":
             lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							1						,							A_  )  )
       elif added_emb_type is None:
             lowerCamelCase_							=   None
       else:
             raise ValueError(
                 f"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""  )

       lowerCamelCase_							=   nn.ModuleList(
           [
               BasicTransformerBlock(
                   A_						,							A_						,							A_						,							dropout=A_						,							activation_fn='gelu'						,							attention_bias=A_						,							)
               for d in range(A_  )
           ]  )

       if norm_in_type == "layer":
             lowerCamelCase_							=   nn.LayerNorm(A_  )
       elif norm_in_type is None:
             lowerCamelCase_							=   None
       else:
             raise ValueError(f"""Unsupported norm_in_type: {norm_in_type}."""  )

       lowerCamelCase_							=   nn.LayerNorm(A_  )

       lowerCamelCase_							=   nn.Linear(A_						,							A_  )

       lowerCamelCase_							=   torch.full(
           [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]						,							-10000.0  )
       causal_attention_mask.triu_(1  )
       lowerCamelCase_							=   causal_attention_mask[None, ...]
       self.register_buffer('causal_attention_mask'						,							A_						,							persistent=A_  )

       lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							A_  )  )
       lowerCamelCase_							=   nn.Parameter(torch.zeros(1						,							A_  )  )

 @property
 # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
 def     a__			(    self :  str  )       ->     Dict[str, AttentionProcessor]:

       """simple docstring"""






       lowerCamelCase_							=   {}

       def fn_recursive_add_processors(A_ :  str						,							A_ :  torch.nn.Module						,							A_ :  Dict[str, AttentionProcessor]  ):
             if hasattr(A_						,							'set_processor'  ):
                   lowerCamelCase_							=   module.processor

             for sub_name, child in module.named_children():
                   fn_recursive_add_processors(f"""{name}.{sub_name}"""						,							A_						,							A_  )

             return processors

       for name, module in self.named_children():
             fn_recursive_add_processors(A_						,							A_						,							A_  )

       return processors

 def     a__			(    self :  List[Any]						,							A_ :  Union[AttentionProcessor, Dict[str, AttentionProcessor]]  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   len(self.attn_processors.keys()  )

       if isinstance(A_						,							A_  ) and len(A_  ) != count:
             raise ValueError(
                 f"""A dict of processors was passed, but the number of processors {len(A_  )} does not match the"""
                 f""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""  )

       def fn_recursive_attn_processor(A_ :  str						,							A_ :  torch.nn.Module						,							A_ :  Union[str, Any]  ):
             if hasattr(A_						,							'set_processor'  ):
                   if not isinstance(A_						,							A_  ):
                         module.set_processor(A_  )
                   else:
                         module.set_processor(processor.pop(f"""{name}.processor"""  )  )

             for sub_name, child in module.named_children():
                   fn_recursive_attn_processor(f"""{name}.{sub_name}"""						,							A_						,							A_  )

       for name, module in self.named_children():
             fn_recursive_attn_processor(A_						,							A_						,							A_  )

 def     a__			(    self :  List[Any]  )       ->     List[Any]:

       """simple docstring"""






       self.set_attn_processor(AttnProcessor()  )

 def     a__			(    self :  Dict						,							A_ :  List[Any]						,							A_ :  Union[torch.Tensor, float, int]						,							A_ :  torch.FloatTensor						,							A_ :  Optional[torch.FloatTensor] = None						,							A_ :  Optional[torch.BoolTensor] = None						,							A_ :  bool = True						,							)       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   hidden_states.shape[0]

       lowerCamelCase_							=   timestep
       if not torch.is_tensor(A_  ):
             lowerCamelCase_							=   torch.tensor([timesteps]						,							dtype=torch.long						,							device=hidden_states.device  )
       elif torch.is_tensor(A_  ) and len(timesteps.shape  ) == 0:
             lowerCamelCase_							=   timesteps[None].to(hidden_states.device  )

       # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
       lowerCamelCase_							=   timesteps * torch.ones(A_						,							dtype=timesteps.dtype						,							device=timesteps.device  )

       lowerCamelCase_							=   self.time_proj(A_  )

       # timesteps does not contain any weights and will always return f32 tensors
       # but time_embedding might be fp16, so we need to cast here.
       lowerCamelCase_							=   timesteps_projected.to(dtype=self.dtype  )
       lowerCamelCase_							=   self.time_embedding(A_  )

       if self.embedding_proj_norm is not None:
             lowerCamelCase_							=   self.embedding_proj_norm(A_  )

       lowerCamelCase_							=   self.embedding_proj(A_  )
       if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
             lowerCamelCase_							=   self.encoder_hidden_states_proj(A_  )
       elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
             raise ValueError('`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'  )

       lowerCamelCase_							=   self.proj_in(A_  )

       lowerCamelCase_							=   self.positional_embedding.to(hidden_states.dtype  )

       lowerCamelCase_							=   []
       lowerCamelCase_							=   0

       if encoder_hidden_states is not None:
             additional_embeds.append(A_  )
             additional_embeddings_len += encoder_hidden_states.shape[1]

       if len(proj_embeddings.shape  ) == 2:
             lowerCamelCase_							=   proj_embeddings[:, None, :]

       if len(hidden_states.shape  ) == 2:
             lowerCamelCase_							=   hidden_states[:, None, :]

       lowerCamelCase_							=   additional_embeds + [
           proj_embeddings,
           time_embeddings[:, None, :],
           hidden_states,
       ]

       if self.prd_embedding is not None:
             lowerCamelCase_							=   self.prd_embedding.to(hidden_states.dtype  ).expand(A_						,							-1						,							-1  )
             additional_embeds.append(A_  )

       lowerCamelCase_							=   torch.cat(
           A_						,							dim=1						,							)

       # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
       lowerCamelCase_							=   additional_embeddings_len + proj_embeddings.shape[1] + 1
       if positional_embeddings.shape[1] < hidden_states.shape[1]:
             lowerCamelCase_							=   F.pad(
                 A_						,							(
                     0,
                     0,
                     additional_embeddings_len,
                     self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
                 )						,							value=0.0						,							)

       lowerCamelCase_							=   hidden_states + positional_embeddings

       if attention_mask is not None:
             lowerCamelCase_							=   (1 - attention_mask.to(hidden_states.dtype  )) * -10000.0
             lowerCamelCase_							=   F.pad(A_						,							(0, self.additional_embeddings)						,							value=0.0  )
             lowerCamelCase_							=   (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype  )
             lowerCamelCase_							=   attention_mask.repeat_interleave(self.config.num_attention_heads						,							dim=0  )

       if self.norm_in is not None:
             lowerCamelCase_							=   self.norm_in(A_  )

       for block in self.transformer_blocks:
             lowerCamelCase_							=   block(A_						,							attention_mask=A_  )

       lowerCamelCase_							=   self.norm_out(A_  )

       if self.prd_embedding is not None:
             lowerCamelCase_							=   hidden_states[:, -1]
       else:
             lowerCamelCase_							=   hidden_states[:, additional_embeddings_len:]

       lowerCamelCase_							=   self.proj_to_clip_embeddings(A_  )

       if not return_dict:
             return (predicted_image_embedding,)

       return PriorTransformerOutput(predicted_image_embedding=A_  )





 def     a__			(    self :  Tuple						,							A_ :  List[Any]  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   (prior_latents * self.clip_std) + self.clip_mean
       return prior_latents



