






'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFRoFormerForCausalLM,
        TFRoFormerForMaskedLM,
        TFRoFormerForMultipleChoice,
        TFRoFormerForQuestionAnswering,
        TFRoFormerForSequenceClassification,
        TFRoFormerForTokenClassification,
        TFRoFormerModel,
    )
    from transformers.models.roformer.modeling_tf_roformer import (
        TFRoFormerSelfAttention,
        TFRoFormerSinusoidalPositionalEmbedding,
    )

class    _UpperCAmelCase			:




   """simple docstring"""

   def __init__(		self :	Tuple      ,      __UpperCAmelCase :	List[Any]      ,      __UpperCAmelCase :	Tuple=13      ,      __UpperCAmelCase :	Optional[int]=7      ,      __UpperCAmelCase :	int=True      ,      __UpperCAmelCase :	str=True      ,      __UpperCAmelCase :	Union[str, Any]=True      ,      __UpperCAmelCase :	str=True      ,      __UpperCAmelCase :	List[str]=99      ,      __UpperCAmelCase :	List[str]=32      ,      __UpperCAmelCase :	Union[str, Any]=2      ,      __UpperCAmelCase :	List[str]=4      ,      __UpperCAmelCase :	Optional[Any]=37      ,      __UpperCAmelCase :	Any="gelu"      ,      __UpperCAmelCase :	Optional[Any]=0.1      ,      __UpperCAmelCase :	int=0.1      ,      __UpperCAmelCase :	Dict=512      ,      __UpperCAmelCase :	List[Any]=16      ,      __UpperCAmelCase :	List[str]=2      ,      __UpperCAmelCase :	Optional[Any]=0.02      ,      __UpperCAmelCase :	int=3      ,      __UpperCAmelCase :	Dict=4      ,      __UpperCAmelCase :	str=None      ,      ):





         '''simple docstring'''





         _A			     =	parent
         _A			     =	13
         _A			     =	7
         _A			     =	True
         _A			     =	True
         _A			     =	True
         _A			     =	True
         _A			     =	99
         _A			     =	32
         _A			     =	2
         _A			     =	4
         _A			     =	37
         _A			     =	"gelu"
         _A			     =	0.1
         _A			     =	0.1
         _A			     =	512
         _A			     =	16
         _A			     =	2
         _A			     =	0.02
         _A			     =	3
         _A			     =	4
         _A			     =	None

   def 						lowerCAmelCase (		self :	Dict		):





         '''simple docstring'''





         _A			     =	ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size		)

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

         _A			     =	None
         if self.use_token_type_ids:
               _A			     =	ids_tensor([self.batch_size, self.seq_length]      ,      self.type_vocab_size		)

         _A			     =	None
         _A			     =	None
         _A			     =	None
         if self.use_labels:
               _A			     =	ids_tensor([self.batch_size]      ,      self.type_sequence_label_size		)
               _A			     =	ids_tensor([self.batch_size, self.seq_length]      ,      self.num_labels		)
               _A			     =	ids_tensor([self.batch_size]      ,      self.num_choices		)

         _A			     =	RoFormerConfig(
             vocab_size=self.vocab_size      ,      hidden_size=self.hidden_size      ,      num_hidden_layers=self.num_hidden_layers      ,      num_attention_heads=self.num_attention_heads      ,      intermediate_size=self.intermediate_size      ,      hidden_act=self.hidden_act      ,      hidden_dropout_prob=self.hidden_dropout_prob      ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,      max_position_embeddings=self.max_position_embeddings      ,      type_vocab_size=self.type_vocab_size      ,      initializer_range=self.initializer_range      ,      return_dict=__UpperCAmelCase      ,      )

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

   def 						lowerCAmelCase (		self :	Optional[Any]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	Union[str, Any]		):





         '''simple docstring'''





         _A			     =	TFRoFormerModel(config=__UpperCAmelCase		)
         _A			     =	{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

         _A			     =	[input_ids, input_mask]
         _A			     =	model(__UpperCAmelCase		)

         _A			     =	model(__UpperCAmelCase		)

         self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)		)

   def 						lowerCAmelCase (		self :	int      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	int      ,      __UpperCAmelCase :	Tuple      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	Tuple      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	List[Any]		):





         '''simple docstring'''





         _A			     =	True
         _A			     =	TFRoFormerForCausalLM(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         _A			     =	model(__UpperCAmelCase		)["logits"]
         self.parent.assertListEqual(
             list(prediction_scores.numpy().shape		)      ,      [self.batch_size, self.seq_length, self.vocab_size]		)

   def 						lowerCAmelCase (		self :	int      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	Tuple      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Tuple      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	str		):





         '''simple docstring'''





         _A			     =	TFRoFormerForMaskedLM(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)

   def 						lowerCAmelCase (		self :	List[Any]      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	List[Any]      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	Tuple      ,      __UpperCAmelCase :	Union[str, Any]		):





         '''simple docstring'''





         _A			     =	self.num_labels
         _A			     =	TFRoFormerForSequenceClassification(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_labels)		)

   def 						lowerCAmelCase (		self :	List[Any]      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	str      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	List[Any]		):





         '''simple docstring'''





         _A			     =	self.num_choices
         _A			     =	TFRoFormerForMultipleChoice(config=__UpperCAmelCase		)
         _A			     =	tf.tile(tf.expand_dims(__UpperCAmelCase      ,      1		)      ,      (1, self.num_choices, 1)		)
         _A			     =	tf.tile(tf.expand_dims(__UpperCAmelCase      ,      1		)      ,      (1, self.num_choices, 1)		)
         _A			     =	tf.tile(tf.expand_dims(__UpperCAmelCase      ,      1		)      ,      (1, self.num_choices, 1)		)
         _A			     =	{
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
             "token_type_ids": multiple_choice_token_type_ids,
         }
         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_choices)		)

   def 						lowerCAmelCase (		self :	Dict      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	int      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	Optional[int]      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Optional[int]		):





         '''simple docstring'''





         _A			     =	self.num_labels
         _A			     =	TFRoFormerForTokenClassification(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.num_labels)		)

   def 						lowerCAmelCase (		self :	List[Any]      ,      __UpperCAmelCase :	List[str]      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	int      ,      __UpperCAmelCase :	Optional[Any]      ,      __UpperCAmelCase :	Any      ,      __UpperCAmelCase :	int      ,      __UpperCAmelCase :	int		):





         '''simple docstring'''





         _A			     =	TFRoFormerForQuestionAnswering(config=__UpperCAmelCase		)
         _A			     =	{
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         _A			     =	model(__UpperCAmelCase		)
         self.parent.assertEqual(result.start_logits.shape      ,      (self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape      ,      (self.batch_size, self.seq_length)		)

   def 						lowerCAmelCase (		self :	Union[str, Any]		):





         '''simple docstring'''





         _A			     =	self.prepare_config_and_inputs()
         (
             (
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	(
             _A
         )     ,	
         )			     =	config_and_inputs
         _A			     =	{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_tf
class    _UpperCAmelCase			(     snake_case_							,    snake_case_							,    unittest.TestCase						):




   """simple docstring"""

   snake_case    				=						(
       (
           TFRoFormerModel,
           TFRoFormerForCausalLM,
           TFRoFormerForMaskedLM,
           TFRoFormerForQuestionAnswering,
           TFRoFormerForSequenceClassification,
           TFRoFormerForTokenClassification,
           TFRoFormerForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   snake_case    				=						(
       {
           '''feature-extraction''': TFRoFormerModel,
           '''fill-mask''': TFRoFormerForMaskedLM,
           '''question-answering''': TFRoFormerForQuestionAnswering,
           '''text-classification''': TFRoFormerForSequenceClassification,
           '''text-generation''': TFRoFormerForCausalLM,
           '''token-classification''': TFRoFormerForTokenClassification,
           '''zero-shot''': TFRoFormerForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )

   snake_case    				=						False
   snake_case    				=						False

   def 						lowerCAmelCase (		self :	int      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	Union[str, Any]      ,      __UpperCAmelCase :	Dict      ,      __UpperCAmelCase :	Optional[int]		):





         '''simple docstring'''





         if pipeline_test_casse_name == "TextGenerationPipelineTests":
               return True

         return False

   def 						lowerCAmelCase (		self :	Optional[int]		):





         '''simple docstring'''





         _A			     =	TFRoFormerModelTester(self		)
         _A			     =	ConfigTester(self      ,      config_class=__UpperCAmelCase      ,      hidden_size=37		)

   def 						lowerCAmelCase (		self :	Optional[int]		):





         '''simple docstring'''





         self.config_tester.run_common_tests()

   def 						lowerCAmelCase (		self :	Optional[int]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Any		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Union[str, Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Optional[Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	List[Any]		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	str		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Any		):





         '''simple docstring'''





         _A			     =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase		)

   @slow
   def 						lowerCAmelCase (		self :	Dict		):





         '''simple docstring'''





         _A			     =	TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base"		)
         self.assertIsNotNone(__UpperCAmelCase		)





@require_tf
class    _UpperCAmelCase			(     unittest.TestCase						):




   """simple docstring"""

   @slow
   def 						lowerCAmelCase (		self :	Optional[int]		):





         '''simple docstring'''





         _A			     =	TFRoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base"		)
         _A			     =	tf.constant([[0, 1, 2, 3, 4, 5]]		)
         _A			     =	model(__UpperCAmelCase		)[0]

         # TODO Replace vocab size
         _A			     =	50000

         _A			     =	[1, 6, vocab_size]
         self.assertEqual(output.shape      ,      __UpperCAmelCase		)

         print(output[:, :3, :3]		)

         # TODO Replace values below with what was printed above.
         _A			     =	tf.constant(
             [
                 [
                     [-0.12053341, -1.0264901, 0.29221946],
                     [-1.5133783, 0.197433, 0.15190607],
                     [-5.0135403, -3.900256, -0.84038764],
                 ]
             ]		)
         tf.debugging.assert_near(output[:, :3, :3]      ,      __UpperCAmelCase      ,      atol=1E-4		)





@require_tf
class    _UpperCAmelCase			(     unittest.TestCase						):




   """simple docstring"""

   snake_case    				=						1E-4

   def 						lowerCAmelCase (		self :	List[str]		):





         '''simple docstring'''





         _A			     =	tf.constant([[4, 10]]		)
         _A			     =	TFRoFormerSinusoidalPositionalEmbedding(num_positions=6      ,      embedding_dim=6		)

         _A			     =	emba(input_ids.shape		)
         _A			     =	tf.constant(
             [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]]		)

         tf.debugging.assert_near(__UpperCAmelCase      ,      __UpperCAmelCase      ,      atol=self.tolerance		)

   def 						lowerCAmelCase (		self :	Tuple		):





         '''simple docstring'''





         _A			     =	tf.constant(
             [
                 [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                 [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
                 [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
             ]		)
         _A			     =	TFRoFormerSinusoidalPositionalEmbedding(num_positions=512      ,      embedding_dim=512		)
         emba([2, 16, 512]		)
         _A			     =	emba.weight[:3, :5]

         tf.debugging.assert_near(__UpperCAmelCase      ,      __UpperCAmelCase      ,      atol=self.tolerance		)





@require_tf
class    _UpperCAmelCase			(     unittest.TestCase						):




   """simple docstring"""

   snake_case    				=						1E-4

   def 						lowerCAmelCase (		self :	str		):





         '''simple docstring'''





         _A			     =	tf.reshape(tf.range(2 * 12 * 16 * 64      ,      dtype=tf.floataa		)      ,      shape=(2, 12, 16, 64)		) / 100

         _A			     =	-tf.reshape(tf.range(2 * 12 * 16 * 64      ,      dtype=tf.floataa		)      ,      shape=(2, 12, 16, 64)		) / 100

         _A			     =	TFRoFormerSinusoidalPositionalEmbedding(num_positions=32      ,      embedding_dim=64		)
         _A			     =	embed_positions([2, 16, 768]		)[None, None, :, :]

         _A     ,	_A			     =	TFRoFormerSelfAttention.apply_rotary_position_embeddings(
             __UpperCAmelCase      ,      __UpperCAmelCase      ,      __UpperCAmelCase		)

         _A			     =	tf.constant(
             [
                 [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
                 [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
                 [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
                 [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
                 [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
                 [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
             ]		)
         _A			     =	tf.constant(
             [
                 [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
                 [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
                 [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
                 [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
                 [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
                 [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
             ]		)

         tf.debugging.assert_near(query_layer[0, 0, :6, :8]      ,      __UpperCAmelCase      ,      atol=self.tolerance		)
         tf.debugging.assert_near(key_layer[0, 0, :6, :8]      ,      __UpperCAmelCase      ,      atol=self.tolerance		)









'''simple docstring'''

import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


a_  :				Dict						= logging.get_logger(__name__)





class      snake_case    (    lowercase			):





  """simple docstring"""




  def __init__(    self			,  UpperCamelCase     ):







   """simple docstring"""
   super().__init__()
   lowerCamelCase_      =  nn.ModuleList(UpperCamelCase     )






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase = None			,  UpperCamelCase = None			,  UpperCamelCase = None			,  UpperCamelCase = None			,  UpperCamelCase = False			,  UpperCamelCase = True			,  ):







   """simple docstring"""
   for i, (image, scale, controlnet) in enumerate(zip(UpperCamelCase			,  UpperCamelCase			,  self.nets     )     ):
    lowerCamelCase_ ,lowerCamelCase_      =  controlnet(
        UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  )

    # merge samples
    if i == 0:
     lowerCamelCase_ ,lowerCamelCase_      =  down_samples, mid_sample
    else:
     lowerCamelCase_      =  [
         samples_prev + samples_curr
         for samples_prev, samples_curr in zip(UpperCamelCase			,  UpperCamelCase     )
     ]
     mid_block_res_sample += mid_sample

   return down_block_res_samples, mid_block_res_sample






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase = True			,  UpperCamelCase = None			,  UpperCamelCase = False			,  UpperCamelCase = None			,  ):







   """simple docstring"""
   lowerCamelCase_      =  0
   lowerCamelCase_      =  save_directory
   for controlnet in self.nets:
    controlnet.save_pretrained(
        UpperCamelCase			,  is_main_process=UpperCamelCase			,  save_function=UpperCamelCase			,  safe_serialization=UpperCamelCase			,  variant=UpperCamelCase			,  )

    idx += 1
    lowerCamelCase_      =  model_path_to_save + f'''_{idx}'''







  @classmethod
  def   snake_case		(    cls			,  UpperCamelCase			,  **UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  0
   lowerCamelCase_      =  []

   # load controlnet and append to list until no controlnet directory exists anymore
   # first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
   # second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
   lowerCamelCase_      =  pretrained_model_path
   while os.path.isdir(UpperCamelCase     ):
    lowerCamelCase_      =  ControlNetModel.from_pretrained(UpperCamelCase			,  **UpperCamelCase     )
    controlnets.append(UpperCamelCase     )

    idx += 1
    lowerCamelCase_      =  pretrained_model_path + f'''_{idx}'''

   logger.info(f'''{len(UpperCamelCase     )} controlnets loaded from {pretrained_model_path}.'''     )

   if len(UpperCamelCase     ) == 0:
    raise ValueError(
        f'''No ControlNets found under {os.path.dirname(UpperCamelCase     )}. Expected at least {pretrained_model_path + "_0"}.'''     )

   return cls(UpperCamelCase     )



