



'''simple docstring'''







import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device





def     snake_case__			(lowerCAmelCase_						):
    UpperCamelCase__	     =      {}
    for name in unet.attn_processors.keys():
        UpperCamelCase__	     =      None if name.endswith("attn1.processor"						) else unet.config.cross_attention_dim
        if name.startswith("mid_block"						):
            UpperCamelCase__	     =      unet.config.block_out_channels[-1]
        elif name.startswith("up_blocks"						):
            UpperCamelCase__	     =      int(name[len("up_blocks."						)]						)
            UpperCamelCase__	     =      list(reversed(unet.config.block_out_channels						)						)[block_id]
        elif name.startswith("down_blocks"						):
            UpperCamelCase__	     =      int(name[len("down_blocks."						)]						)
            UpperCamelCase__	     =      unet.config.block_out_channels[block_id]
        UpperCamelCase__	     =      (
            LoRAAttnProcessor2_0 if hasattr(UpperCamelCase__       ,	"scaled_dot_product_attention"						) else LoRAAttnProcessor
        )
        UpperCamelCase__	     =      lora_attn_processor_class(
            hidden_size=UpperCamelCase__       ,	cross_attention_dim=UpperCamelCase__						)
    UpperCamelCase__	     =      AttnProcsLayers(UpperCamelCase__						)
    return lora_attn_procs, unet_lora_layers





def     snake_case__			(lowerCAmelCase_						):
    UpperCamelCase__	     =      {}
    UpperCamelCase__	     =      (
        LoRAAttnProcessor2_0 if hasattr(UpperCamelCase__       ,	"scaled_dot_product_attention"						) else LoRAAttnProcessor
    )
    for name, module in text_encoder.named_modules():
        if name.endswith(UpperCamelCase__						):
            UpperCamelCase__	     =      lora_attn_processor_class(
                hidden_size=module.out_proj.out_features       ,	cross_attention_dim=UpperCamelCase__						)
    return text_lora_attn_procs





def     snake_case__			(lowerCAmelCase_						):
    UpperCamelCase__	     =      create_text_encoder_lora_attn_procs(UpperCamelCase__						)
    UpperCamelCase__	     =      AttnProcsLayers(UpperCamelCase__						)
    return text_encoder_lora_layers





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_=False						):


    for _, attn_proc in text_lora_attn_procs.items():
        # set up.weights
        for layer_name, layer_module in attn_proc.named_modules():
            if layer_name.endswith("_lora"						):
                UpperCamelCase__	     =      (
                    torch.randn_like(layer_module.up.weight						)
                    if randn_weight
                    else torch.zeros_like(layer_module.up.weight						)
                )
                UpperCamelCase__	     =      torch.nn.Parameter(UpperCamelCase__						)





class    lowerCamelCase_				(						unittest.TestCase	):






   '''simple docstring'''

   def 						_UpperCamelCase							(	self			):
       torch.manual_seed(0			)
       UpperCamelCase__	     =      UNetaDConditionModel(
           block_out_channels=(32, 64)  ,    layers_per_block=2  ,    sample_size=32  ,    in_channels=4  ,    out_channels=4  ,    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")  ,    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")  ,    cross_attention_dim=32  ,    )
       UpperCamelCase__	     =      DDIMScheduler(
           beta_start=0.00_085  ,    beta_end=0.012  ,    beta_schedule="scaled_linear"  ,    clip_sample=A_  ,    set_alpha_to_one=A_  ,    steps_offset=1  ,    )
       torch.manual_seed(0			)
       UpperCamelCase__	     =      AutoencoderKL(
           block_out_channels=[32, 64]  ,    in_channels=3  ,    out_channels=3  ,    down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]  ,    up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]  ,    latent_channels=4  ,    )
       UpperCamelCase__	     =      CLIPTextConfig(
           bos_token_id=0  ,    eos_token_id=2  ,    hidden_size=32  ,    intermediate_size=37  ,    layer_norm_eps=1e-05  ,    num_attention_heads=4  ,    num_hidden_layers=5  ,    pad_token_id=1  ,    vocab_size=10_00  ,    )
       UpperCamelCase__	     =      CLIPTextModel(A_			)
       UpperCamelCase__	     =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

       UpperCamelCase__     ,	UpperCamelCase__	     =      create_unet_lora_layers(A_			)
       UpperCamelCase__	     =      create_text_encoder_lora_layers(A_			)

       UpperCamelCase__	     =      {
           "unet": unet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "safety_checker": None,
           "feature_extractor": None,
       }
       UpperCamelCase__	     =      {
           "unet_lora_layers": unet_lora_layers,
           "text_encoder_lora_layers": text_encoder_lora_layers,
           "unet_lora_attn_procs": unet_lora_attn_procs,
       }
       return pipeline_components, lora_components


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      1
       UpperCamelCase__	     =      10
       UpperCamelCase__	     =      4
       UpperCamelCase__	     =      (32, 32)

       UpperCamelCase__	     =      torch.manual_seed(0			)
       UpperCamelCase__	     =      floats_tensor((batch_size, num_channels) + sizes			)
       UpperCamelCase__	     =      torch.randint(1  ,    A_  ,    size=(batch_size, sequence_length)  ,    generator=A_			)

       UpperCamelCase__	     =      {
           "prompt": "A painting of a squirrel eating a burger",
           "generator": generator,
           "num_inference_steps": 2,
           "guidance_scale": 6.0,
           "output_type": "numpy",
       }

       return noise, input_ids, pipeline_inputs

       # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      77

       UpperCamelCase__	     =      torch.randint(2  ,    56  ,    size=(1, max_seq_length)  ,    generator=torch.manual_seed(0			)			)

       UpperCamelCase__	     =      {}
       UpperCamelCase__	     =      inputs
       return prepared_inputs


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       LoraLoaderMixin.save_lora_weights(
           save_directory=A_  ,    unet_lora_layers=lora_components["unet_lora_layers"]  ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]  ,    )
       self.assertTrue(os.path.isfile(os.path.join(A_  ,    "pytorch_lora_weights.bin"			)			)			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
       UpperCamelCase__	     =      sd_pipe.to(A_			)
       sd_pipe.set_progress_bar_config(disable=A_			)

       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_inputs()

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
           LoraLoaderMixin.save_lora_weights(
               save_directory=A_  ,    unet_lora_layers=lora_components["unet_lora_layers"]  ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]  ,    )
           self.assertTrue(os.path.isfile(os.path.join(A_  ,    "pytorch_lora_weights.bin"			)			)			)
           sd_pipe.load_lora_weights(A_			)

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A_			)  ,    torch.from_numpy(A_			)			)			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
       UpperCamelCase__	     =      sd_pipe.to(A_			)
       sd_pipe.set_progress_bar_config(disable=A_			)

       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_inputs()

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
           LoraLoaderMixin.save_lora_weights(
               save_directory=A_  ,    unet_lora_layers=lora_components["unet_lora_layers"]  ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]  ,    safe_serialization=A_  ,    )
           self.assertTrue(os.path.isfile(os.path.join(A_  ,    "pytorch_lora_weights.safetensors"			)			)			)
           sd_pipe.load_lora_weights(A_			)

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A_			)  ,    torch.from_numpy(A_			)			)			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      lora_components["unet_lora_attn_procs"]
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
       UpperCamelCase__	     =      sd_pipe.to(A_			)
       sd_pipe.set_progress_bar_config(disable=A_			)

       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_inputs()

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
           UpperCamelCase__	     =      sd_pipe.unet
           unet.set_attn_processor(A_			)
           unet.save_attn_procs(A_			)
           self.assertTrue(os.path.isfile(os.path.join(A_  ,    "pytorch_lora_weights.bin"			)			)			)
           sd_pipe.load_lora_weights(A_			)

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A_			)  ,    torch.from_numpy(A_			)			)			)


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)

       UpperCamelCase__	     =      self.get_dummy_tokens()

       # inference without lora
       UpperCamelCase__	     =      pipe.text_encoder(**A_			)[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with zeroed out up.weights
       UpperCamelCase__	     =      create_text_encoder_lora_attn_procs(pipe.text_encoder			)
       set_lora_up_weights(A_  ,    randn_weight=A_			)

       # monkey patch
       pipe._modify_text_encoder(A_			)

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       UpperCamelCase__	     =      pipe.text_encoder(**A_			)[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert torch.allclose(
           A_  ,    A_			), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

       # create lora_attn_procs with randn up.weights
       UpperCamelCase__	     =      create_text_encoder_lora_attn_procs(pipe.text_encoder			)
       set_lora_up_weights(A_  ,    randn_weight=A_			)

       # monkey patch
       pipe._modify_text_encoder(A_			)

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       UpperCamelCase__	     =      pipe.text_encoder(**A_			)[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert not torch.allclose(
           A_  ,    A_			), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)

       UpperCamelCase__	     =      self.get_dummy_tokens()

       # inference without lora
       UpperCamelCase__	     =      pipe.text_encoder(**A_			)[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with randn up.weights
       UpperCamelCase__	     =      create_text_encoder_lora_attn_procs(pipe.text_encoder			)
       set_lora_up_weights(A_  ,    randn_weight=A_			)

       # monkey patch
       pipe._modify_text_encoder(A_			)

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       UpperCamelCase__	     =      pipe.text_encoder(**A_			)[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert not torch.allclose(
           A_  ,    A_			), "lora outputs should be different to without lora outputs"

       # remove monkey patch
       pipe._remove_text_encoder_monkey_patch()

       # inference with removed lora
       UpperCamelCase__	     =      pipe.text_encoder(**A_			)[0]
       assert outputs_without_lora_removed.shape == (1, 77, 32)

       assert torch.allclose(
           A_  ,    A_			), "remove lora monkey patch should restore the original outputs"


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
       UpperCamelCase__	     =      sd_pipe.to(A_			)
       sd_pipe.set_progress_bar_config(disable=A_			)

       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_inputs()

       with tempfile.TemporaryDirectory() as tmpdirname:
           LoraLoaderMixin.save_lora_weights(
               save_directory=A_  ,    unet_lora_layers=lora_components["unet_lora_layers"]  ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]  ,    )
           self.assertTrue(os.path.isfile(os.path.join(A_  ,    "pytorch_lora_weights.bin"			)			)			)
           sd_pipe.load_lora_weights(A_			)

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      lora_images[0, -3:, -3:, -1]

       UpperCamelCase__	     =      sd_pipe(**A_  ,    cross_attention_kwargs={"scale": 0.5}			).images
       UpperCamelCase__	     =      lora_images_with_scale[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(
           torch.allclose(torch.from_numpy(A_			)  ,    torch.from_numpy(A_			)			)			)


   def 						_UpperCamelCase							(	self			):



       with tempfile.TemporaryDirectory() as tmpdirname:
           self.create_lora_weight_file(A_			)

           UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
           UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
           UpperCamelCase__	     =      sd_pipe.to(A_			)
           sd_pipe.set_progress_bar_config(disable=A_			)

           # check if vanilla attention processors are used
           for _, module in sd_pipe.unet.named_modules():
               if isinstance(A_  ,    A_			):
                   self.assertIsInstance(module.processor  ,    (AttnProcessor, AttnProcessor2_0)			)

                   # load LoRA weight file
           sd_pipe.load_lora_weights(A_			)

           # check if lora attention processors are used
           for _, module in sd_pipe.unet.named_modules():
               if isinstance(A_  ,    A_			):
                   UpperCamelCase__	     =      (
                       LoRAAttnProcessor2_0 if hasattr(A_  ,    "scaled_dot_product_attention"			) else LoRAAttnProcessor
                   )
                   self.assertIsInstance(module.processor  ,    A_			)


   @unittest.skipIf(torch_device != "cuda"  ,    "This test is supposed to run on GPU"			)
   def 						_UpperCamelCase							(	self			):



       with tempfile.TemporaryDirectory() as tmpdirname:
           self.create_lora_weight_file(A_			)

           UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
           UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
           UpperCamelCase__	     =      sd_pipe.to(A_			)
           sd_pipe.set_progress_bar_config(disable=A_			)

           # enable XFormers
           sd_pipe.enable_xformers_memory_efficient_attention()

           # check if xFormers attention processors are used
           for _, module in sd_pipe.unet.named_modules():
               if isinstance(A_  ,    A_			):
                   self.assertIsInstance(module.processor  ,    A_			)

                   # load LoRA weight file
           sd_pipe.load_lora_weights(A_			)

           # check if lora attention processors are used
           for _, module in sd_pipe.unet.named_modules():
               if isinstance(A_  ,    A_			):
                   self.assertIsInstance(module.processor  ,    A_			)







   @unittest.skipIf(torch_device != "cuda"  ,    "This test is supposed to run on GPU"			)
   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_components()
       UpperCamelCase__	     =      StableDiffusionPipeline(**A_			)
       UpperCamelCase__	     =      sd_pipe.to(A_			)
       sd_pipe.set_progress_bar_config(disable=A_			)

       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      self.get_dummy_inputs()

       # enable XFormers
       sd_pipe.enable_xformers_memory_efficient_attention()

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
           LoraLoaderMixin.save_lora_weights(
               save_directory=A_  ,    unet_lora_layers=lora_components["unet_lora_layers"]  ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]  ,    )
           self.assertTrue(os.path.isfile(os.path.join(A_  ,    "pytorch_lora_weights.bin"			)			)			)
           sd_pipe.load_lora_weights(A_			)

       UpperCamelCase__	     =      sd_pipe(**A_			).images
       UpperCamelCase__	     =      lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A_			)  ,    torch.from_numpy(A_			)			)			)









"""simple docstring"""






import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
        RwkvForCausalLM,
        RwkvModel,
    )
    from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
    lowerCamelCase				:     str								=  False




class    __snake_case:




    def __init__(   self    ,							A_    ,							A_=14    ,							A_=7    ,							A_=True    ,							A_=False    ,							A_=True    ,							A_=True    ,							A_=True    ,							A_=99    ,							A_=32    ,							A_=5    ,							A_=37    ,							A_="gelu"    ,							A_=0.1    ,							A_=0.1    ,							A_=512    ,							A_=16    ,							A_=2    ,							A_=3    ,							A_=4    ,							A_=None    ,							):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			parent
      _SCREAMING_SNAKE_CASE        =			batch_size
      _SCREAMING_SNAKE_CASE        =			seq_length
      _SCREAMING_SNAKE_CASE        =			is_training
      _SCREAMING_SNAKE_CASE        =			use_token_type_ids
      _SCREAMING_SNAKE_CASE        =			use_input_mask
      _SCREAMING_SNAKE_CASE        =			use_labels
      _SCREAMING_SNAKE_CASE        =			use_mc_token_ids
      _SCREAMING_SNAKE_CASE        =			vocab_size
      _SCREAMING_SNAKE_CASE        =			hidden_size
      _SCREAMING_SNAKE_CASE        =			num_hidden_layers
      _SCREAMING_SNAKE_CASE        =			intermediate_size
      _SCREAMING_SNAKE_CASE        =			hidden_act
      _SCREAMING_SNAKE_CASE        =			hidden_dropout_prob
      _SCREAMING_SNAKE_CASE        =			attention_probs_dropout_prob
      _SCREAMING_SNAKE_CASE        =			max_position_embeddings
      _SCREAMING_SNAKE_CASE        =			type_vocab_size
      _SCREAMING_SNAKE_CASE        =			type_sequence_label_size
      _SCREAMING_SNAKE_CASE        =			num_labels
      _SCREAMING_SNAKE_CASE        =			num_choices
      _SCREAMING_SNAKE_CASE        =			scope
      _SCREAMING_SNAKE_CASE        =			vocab_size - 1
      _SCREAMING_SNAKE_CASE        =			vocab_size - 1
      _SCREAMING_SNAKE_CASE        =			vocab_size - 1




    def   A   (   self			):





      '''simple docstring'''



      return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''			)




    def   A   (   self    ,							A_=False    ,							A_=False    ,							A_=False			):





      '''simple docstring'''



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

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

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

      _SCREAMING_SNAKE_CASE        =			None
      if self.use_mc_token_ids:
        _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size, self.num_choices]    ,							self.seq_length			)

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

      _SCREAMING_SNAKE_CASE        =			self.get_config(
          gradient_checkpointing=A_    ,							scale_attn_by_inverse_layer_idx=A_    ,							reorder_and_upcast_attn=A_    ,							)

      return (
          config,
          input_ids,
          input_mask,
          None,
          token_type_ids,
          mc_token_ids,
          sequence_labels,
          token_labels,
          choice_labels,
      )




    def   A   (   self    ,							A_=False    ,							A_=False    ,							A_=False			):





      '''simple docstring'''



      return RwkvConfig(
          vocab_size=self.vocab_size    ,							hidden_size=self.hidden_size    ,							num_hidden_layers=self.num_hidden_layers    ,							intermediate_size=self.intermediate_size    ,							activation_function=self.hidden_act    ,							resid_pdrop=self.hidden_dropout_prob    ,							attn_pdrop=self.attention_probs_dropout_prob    ,							n_positions=self.max_position_embeddings    ,							type_vocab_size=self.type_vocab_size    ,							use_cache=A_    ,							bos_token_id=self.bos_token_id    ,							eos_token_id=self.eos_token_id    ,							pad_token_id=self.pad_token_id    ,							gradient_checkpointing=A_    ,							scale_attn_by_inverse_layer_idx=A_    ,							reorder_and_upcast_attn=A_    ,							)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.get_config()
      _SCREAMING_SNAKE_CASE        =			300
      return config




    def   A   (   self			):





      '''simple docstring'''



      (
          (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     
      )        =			self.prepare_config_and_inputs()

      _SCREAMING_SNAKE_CASE        =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
      _SCREAMING_SNAKE_CASE        =			ids_tensor([self.batch_size, self.seq_length]    ,							vocab_size=2			)

      return (
          config,
          input_ids,
          input_mask,
          head_mask,
          token_type_ids,
          sequence_labels,
          token_labels,
          choice_labels,
          encoder_hidden_states,
          encoder_attention_mask,
      )




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							*A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			True
      _SCREAMING_SNAKE_CASE        =			RwkvModel(config=A_			)
      model.to(A_			)
      model.eval()

      _SCREAMING_SNAKE_CASE        =			model(A_			)

      self.parent.assertEqual(result.last_hidden_state.shape    ,							(self.batch_size, self.seq_length, self.hidden_size)			)
      self.parent.assertEqual(len(result.hidden_states			)    ,							config.num_hidden_layers + 1			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							*A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			RwkvForCausalLM(A_			)
      model.to(A_			)
      model.eval()

      _SCREAMING_SNAKE_CASE        =			model(A_    ,							labels=A_			)
      self.parent.assertEqual(result.loss.shape    ,							()			)
      self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.vocab_size)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							*A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			RwkvModel(config=A_			)
      model.to(A_			)
      model.eval()

      _SCREAMING_SNAKE_CASE        =			model(A_			)
      _SCREAMING_SNAKE_CASE        =			outputs.last_hidden_state

      _SCREAMING_SNAKE_CASE        =			model(input_ids[:, :2]			)
      _SCREAMING_SNAKE_CASE        =			outputs.last_hidden_state

      # Using the state computed on the first inputs, we will get the same output
      _SCREAMING_SNAKE_CASE        =			model(input_ids[:, 2:]    ,							state=outputs.state			)
      _SCREAMING_SNAKE_CASE        =			outputs.last_hidden_state

      self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]    ,							dim=1			)    ,							A_    ,							atol=1e-5			)			)




    def   A   (   self    ,							A_    ,							A_    ,							A_    ,							A_    ,							A_    ,							*A_    ,							A_=False			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			RwkvForCausalLM(A_			)
      model.to(A_			)
      if gradient_checkpointing:
        model.gradient_checkpointing_enable()

      _SCREAMING_SNAKE_CASE        =			model(A_    ,							labels=A_			)
      self.parent.assertEqual(result.loss.shape    ,							()			)
      self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.vocab_size)			)
      result.loss.backward()




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.prepare_config_and_inputs()

      (
          (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     (
          _SCREAMING_SNAKE_CASE
      )    ,     
      )        =			config_and_inputs

      _SCREAMING_SNAKE_CASE        =			{'''input_ids''': input_ids}

      return config, inputs_dict




@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0			,			reason='''See https://github.com/huggingface/transformers/pull/24204'''       )
@require_torch
class    __snake_case(					__A			,			__A			,			__A			,			unittest.TestCase       ):
    _A       = (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
    _A       = (
        {'''feature-extraction''': RwkvModel, '''text-generation''': RwkvForCausalLM} if is_torch_available() else {}
    )
    # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
    _A       = False
    _A       = False
    _A       = False
    _A       = False
    _A       = False  # Rwkv does not support head masking




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			RwkvModelTester(self			)
      _SCREAMING_SNAKE_CASE        =			ConfigTester(
          self    ,							config_class=A_    ,							n_embd=37    ,							common_properties=['''hidden_size''', '''num_hidden_layers''']			)




    def   A   (   self    ,							A_    ,							A_    ,							A_=None			):





      '''simple docstring'''



      if isinstance(A_    ,							torch.Tensor			):
        _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			member.max().item(), member.min().item()
      elif isinstance(A_    ,							A_			) or isinstance(A_    ,							A_			):
        _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			max(A_			), min(A_			)

      if not isinstance(A_    ,							A_			):
        raise TypeError('''container should be a list or tuple'''			)
      elif len(A_			) != 2:
        raise ValueError('''container should have 2 elements'''			)

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			container

      _SCREAMING_SNAKE_CASE        =			(min_value >= expected_min) and (max_value <= expected_max)

      if not is_inside_interval:
        _SCREAMING_SNAKE_CASE        =			'''%s not found in %s''' % (safe_repr(A_			), safe_repr(A_			))
        self.fail(self._formatMessage(A_    ,							A_			)			)




    def   A   (   self			):





      '''simple docstring'''



      self.config_tester.run_common_tests()




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_rwkv_model(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_causl_lm(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_state_equivalency(*A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        _SCREAMING_SNAKE_CASE        =			model_class(config=A_			)
        for name, param in model.named_parameters():
          if "time_decay" in name:
            if param.requires_grad:
              self.assertTrue(param.data.max().item() == 3.0			)
              self.assertTrue(param.data.min().item() == -5.0			)
          elif "time_first" in name:
            if param.requires_grad:
              # check if it's a ones like
              self.assertTrue(torch.allclose(param.data    ,							torch.ones_like(param.data			)    ,							atol=1e-5    ,							rtol=1e-5			)			)
          elif any(x in name for x in ['''time_mix_key''', '''time_mix_receptance''']			):
            if param.requires_grad:
              self.assertInterval(
                  param.data    ,							[0.0, 1.0]    ,							msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''    ,							)
          elif "time_mix_value" in name:
            if param.requires_grad:
              self.assertInterval(
                  param.data    ,							[0.0, 1.3]    ,							msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''    ,							)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.model_tester.prepare_config_and_inputs_for_common()
      _SCREAMING_SNAKE_CASE        =			True

      _SCREAMING_SNAKE_CASE        =			getattr(self.model_tester    ,							'''seq_length'''    ,							A_			)

      for model_class in self.all_model_classes:
        _SCREAMING_SNAKE_CASE        =			True
        _SCREAMING_SNAKE_CASE        =			False
        _SCREAMING_SNAKE_CASE        =			True
        _SCREAMING_SNAKE_CASE        =			model_class(A_			)
        model.to(A_			)
        model.eval()

        _SCREAMING_SNAKE_CASE        =			self._prepare_for_class(A_    ,							A_			)
        _SCREAMING_SNAKE_CASE        =			inputs['''input_ids'''].shape[0]
        with torch.no_grad():
          _SCREAMING_SNAKE_CASE        =			model(**A_			)
        _SCREAMING_SNAKE_CASE        =			outputs.attentions
        self.assertEqual(len(A_			)    ,							self.model_tester.num_hidden_layers			)

        # check that output_attentions also work using config
        del inputs_dict["output_attentions"]
        _SCREAMING_SNAKE_CASE        =			True
        _SCREAMING_SNAKE_CASE        =			model_class(A_			)
        model.to(A_			)
        model.eval()

        _SCREAMING_SNAKE_CASE        =			self._prepare_for_class(A_    ,							A_			)
        _SCREAMING_SNAKE_CASE        =			inputs['''input_ids'''].shape[0]
        with torch.no_grad():
          _SCREAMING_SNAKE_CASE        =			model(**A_			)
        _SCREAMING_SNAKE_CASE        =			outputs.attentions
        self.assertEqual(len(A_			)    ,							self.model_tester.num_hidden_layers			)

        self.assertListEqual(
            list(attentions[0].shape[-3:]			)    ,							[batch_size, seq_len, config.hidden_size]    ,							)
        _SCREAMING_SNAKE_CASE        =			len(A_			)

        # Check attention is always last and order is fine
        _SCREAMING_SNAKE_CASE        =			True
        _SCREAMING_SNAKE_CASE        =			True
        _SCREAMING_SNAKE_CASE        =			model_class(A_			)
        model.to(A_			)
        model.eval()

        _SCREAMING_SNAKE_CASE        =			self._prepare_for_class(A_    ,							A_			)
        _SCREAMING_SNAKE_CASE        =			inputs['''input_ids'''].shape[0]
        with torch.no_grad():
          _SCREAMING_SNAKE_CASE        =			model(**A_			)

        _SCREAMING_SNAKE_CASE        =			1
        self.assertEqual(out_len + added_hidden_states    ,							len(A_			)			)

        _SCREAMING_SNAKE_CASE        =			outputs.attentions

        self.assertEqual(len(A_			)    ,							self.model_tester.num_hidden_layers			)
        self.assertListEqual(
            list(self_attentions[0].shape[-3:]			)    ,							[batch_size, seq_len, config.hidden_size]    ,							)




    @slow
    def   A   (   self			):





      '''simple docstring'''



      for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _SCREAMING_SNAKE_CASE        =			RwkvModel.from_pretrained(A_			)
        self.assertIsNotNone(A_			)




@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0			,			reason='''See https://github.com/huggingface/transformers/pull/24204'''       )
@slow
class    __snake_case(					unittest.TestCase       ):




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			'''RWKV/rwkv-4-169m-pile'''
      _SCREAMING_SNAKE_CASE        =			AutoTokenizer.from_pretrained(self.model_id			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			'''Hello my name is Jasmine and I am a newbie to the'''
      _SCREAMING_SNAKE_CASE        =			RwkvForCausalLM.from_pretrained(self.model_id			).to(A_			)

      _SCREAMING_SNAKE_CASE        =			self.tokenizer('''Hello my name is'''    ,							return_tensors='''pt'''			).input_ids.to(A_			)
      _SCREAMING_SNAKE_CASE        =			model.generate(A_    ,							max_new_tokens=10			)
      _SCREAMING_SNAKE_CASE        =			self.tokenizer.decode(output[0].tolist()			)

      self.assertEqual(A_    ,							A_			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			'''Hello my name is Jasmine and I am a newbie to the'''

      _SCREAMING_SNAKE_CASE        =			self.tokenizer('''Hello my name is'''    ,							return_tensors='''pt'''			).input_ids.to(A_			)
      _SCREAMING_SNAKE_CASE        =			RwkvForCausalLM.from_pretrained(self.model_id    ,							torch_dtype=torch.bfloataa			).to(A_			)

      _SCREAMING_SNAKE_CASE        =			model.generate(A_    ,							max_new_tokens=10			)
      _SCREAMING_SNAKE_CASE        =			self.tokenizer.decode(output[0].tolist()			)

      self.assertEqual(A_    ,							A_			)
