
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:
						snake_case__  							=							False







class 							UpperCAmelCase	:

					def __init__(			self      :					int ,      lowerCAmelCase      :					Union[str, Any] ,      lowerCAmelCase      :					Any=14 ,      lowerCAmelCase      :					Optional[Any]=7 ,      lowerCAmelCase      :					List[Any]=True ,      lowerCAmelCase      :					str=False ,      lowerCAmelCase      :					Any=True ,      lowerCAmelCase      :					int=True ,      lowerCAmelCase      :					str=True ,      lowerCAmelCase      :					List[Any]=99 ,      lowerCAmelCase      :					str=32 ,      lowerCAmelCase      :					Union[str, Any]=5 ,      lowerCAmelCase      :					str=37 ,      lowerCAmelCase      :					Dict="gelu" ,      lowerCAmelCase      :					List[Any]=0.1 ,      lowerCAmelCase      :					Union[str, Any]=0.1 ,      lowerCAmelCase      :					str=512 ,      lowerCAmelCase      :					Dict=16 ,      lowerCAmelCase      :					str=2 ,      lowerCAmelCase      :					Dict=3 ,      lowerCAmelCase      :					List[Any]=4 ,      lowerCAmelCase      :					int=None ,      ):
											lowercase       :    str											=  parent
											lowercase       :    List[Any]											=  batch_size
											lowercase       :    Any											=  seq_length
											lowercase       :    Tuple											=  is_training
											lowercase       :    List[Any]											=  use_token_type_ids
											lowercase       :    Any											=  use_input_mask
											lowercase       :    Dict											=  use_labels
											lowercase       :    Union[str, Any]											=  use_mc_token_ids
											lowercase       :    List[Any]											=  vocab_size
											lowercase       :    List[str]											=  hidden_size
											lowercase       :    int											=  num_hidden_layers
											lowercase       :    Dict											=  intermediate_size
											lowercase       :    Dict											=  hidden_act
											lowercase       :    List[str]											=  hidden_dropout_prob
											lowercase       :    Tuple											=  attention_probs_dropout_prob
											lowercase       :    Optional[Any]											=  max_position_embeddings
											lowercase       :    Optional[int]											=  type_vocab_size
											lowercase       :    Tuple											=  type_sequence_label_size
											lowercase       :    str											=  num_labels
											lowercase       :    Optional[int]											=  num_choices
											lowercase       :    int											=  scope
											lowercase       :    Optional[Any]											=  vocab_size - 1
											lowercase       :    List[str]											=  vocab_size - 1
											lowercase       :    Union[str, Any]											=  vocab_size - 1

					def        _lowerCAmelCase	(			self      :					Union[str, Any]					):
											return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''					)

					def        _lowerCAmelCase	(			self      :					Optional[Any] ,      lowerCAmelCase      :					Union[str, Any]=False ,      lowerCAmelCase      :					int=False ,      lowerCAmelCase      :					str=False					):
											lowercase       :    Dict											=  ids_tensor([self.batch_size, self.seq_length] ,      self.vocab_size					)

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

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

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

											lowercase       :    str											=  None
											lowercase       :    str											=  None
											lowercase       :    int											=  None
											if self.use_labels:
																	lowercase       :    str											=  ids_tensor([self.batch_size] ,      self.type_sequence_label_size					)
																	lowercase       :    Union[str, Any]											=  ids_tensor([self.batch_size, self.seq_length] ,      self.num_labels					)
																	lowercase       :    Dict											=  ids_tensor([self.batch_size] ,      self.num_choices					)

											lowercase       :    Tuple											=  self.get_config(
											    gradient_checkpointing=lowerCAmelCase ,      scale_attn_by_inverse_layer_idx=lowerCAmelCase ,      reorder_and_upcast_attn=lowerCAmelCase ,      )

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

					def        _lowerCAmelCase	(			self      :					Optional[Any] ,      lowerCAmelCase      :					Any=False ,      lowerCAmelCase      :					Dict=False ,      lowerCAmelCase      :					List[str]=False					):
											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=lowerCAmelCase ,      bos_token_id=self.bos_token_id ,      eos_token_id=self.eos_token_id ,      pad_token_id=self.pad_token_id ,      gradient_checkpointing=lowerCAmelCase ,      scale_attn_by_inverse_layer_idx=lowerCAmelCase ,      reorder_and_upcast_attn=lowerCAmelCase ,      )

					def        _lowerCAmelCase	(			self      :					Any					):
											lowercase       :    Dict											=  self.get_config()
											lowercase       :    Tuple											=  300
											return config

					def        _lowerCAmelCase	(			self      :					Union[str, Any]					):
											(
											    lowercase				
											)       :    Tuple											=  self.prepare_config_and_inputs()

											lowercase       :    Optional[Any]											=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
											lowercase       :    Any											=  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        _lowerCAmelCase	(			self      :					List[Any] ,      lowerCAmelCase      :					Optional[Any] ,      lowerCAmelCase      :					str ,      lowerCAmelCase      :					int ,      lowerCAmelCase      :					List[str] ,      lowerCAmelCase      :					Optional[int] ,      *lowerCAmelCase      :					Tuple					):
											lowercase       :    Optional[Any]											=  True
											lowercase       :    List[Any]											=  RwkvModel(config=lowerCAmelCase					)
											model.to(lowerCAmelCase					)
											model.eval()

											lowercase       :    Dict											=  model(lowerCAmelCase					)

											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        _lowerCAmelCase	(			self      :					Union[str, Any] ,      lowerCAmelCase      :					Optional[Any] ,      lowerCAmelCase      :					Any ,      lowerCAmelCase      :					str ,      lowerCAmelCase      :					Optional[Any] ,      lowerCAmelCase      :					Optional[Any] ,      *lowerCAmelCase      :					Dict					):
											lowercase       :    str											=  RwkvForCausalLM(lowerCAmelCase					)
											model.to(lowerCAmelCase					)
											model.eval()

											lowercase       :    List[str]											=  model(lowerCAmelCase ,      labels=lowerCAmelCase					)
											self.parent.assertEqual(result.loss.shape ,      ()					)
											self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.vocab_size)					)

					def        _lowerCAmelCase	(			self      :					Tuple ,      lowerCAmelCase      :					Optional[Any] ,      lowerCAmelCase      :					List[str] ,      lowerCAmelCase      :					Any ,      lowerCAmelCase      :					Optional[int] ,      lowerCAmelCase      :					str ,      *lowerCAmelCase      :					Dict					):
											lowercase       :    Any											=  RwkvModel(config=lowerCAmelCase					)
											model.to(lowerCAmelCase					)
											model.eval()

											lowercase       :    Optional[int]											=  model(lowerCAmelCase					)
											lowercase       :    int											=  outputs.last_hidden_state

											lowercase       :    int											=  model(input_ids[:, :2]					)
											lowercase       :    Optional[int]											=  outputs.last_hidden_state

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

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

					def        _lowerCAmelCase	(			self      :					Tuple ,      lowerCAmelCase      :					List[str] ,      lowerCAmelCase      :					List[str] ,      lowerCAmelCase      :					Optional[Any] ,      lowerCAmelCase      :					Any ,      lowerCAmelCase      :					Optional[Any] ,      *lowerCAmelCase      :					List[str] ,      lowerCAmelCase      :					Any=False					):
											lowercase       :    Optional[Any]											=  RwkvForCausalLM(lowerCAmelCase					)
											model.to(lowerCAmelCase					)
											if gradient_checkpointing:
																	model.gradient_checkpointing_enable()

											lowercase       :    Any											=  model(lowerCAmelCase ,      labels=lowerCAmelCase					)
											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        _lowerCAmelCase	(			self      :					int					):
											lowercase       :    Optional[Any]											=  self.prepare_config_and_inputs()

											(
											    lowercase				
											)       :    Optional[Any]											=  config_and_inputs

											lowercase       :    List[str]											=  {'''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 							UpperCAmelCase	(   __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, unittest.TestCase					):
					a__: Dict			  =				(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
					a__: Dict			  =				(
					    {"""feature-extraction""": RwkvModel, """text-generation""": RwkvForCausalLM} if is_torch_available() else {}
					)
					# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
					a__: str			  =				False
					a__: int			  =				False
					a__: List[Any]			  =				False
					a__: str			  =				False
					a__: Any			  =				False  # Rwkv does not support head masking

					def        _lowerCAmelCase	(			self      :					Any					):
											lowercase       :    List[str]											=  RwkvModelTester(self					)
											lowercase       :    Optional[Any]											=  ConfigTester(
											    self ,      config_class=lowerCAmelCase ,      n_embd=37 ,      common_properties=['''hidden_size''', '''num_hidden_layers''']					)

					def        _lowerCAmelCase	(			self      :					Optional[Any] ,      lowerCAmelCase      :					Optional[int] ,      lowerCAmelCase      :					List[Any] ,      lowerCAmelCase      :					Union[str, Any]=None					):
											if isinstance(lowerCAmelCase ,      torch.Tensor					):
																	lowercase       :    List[str]											=  member.max().item(), member.min().item()
											elif isinstance(lowerCAmelCase ,      lowerCAmelCase					) or isinstance(lowerCAmelCase ,      lowerCAmelCase					):
																	lowercase       :    Optional[Any]											=  max(lowerCAmelCase					), min(lowerCAmelCase					)

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

											lowercase       :    Dict											=  container

											lowercase       :    List[Any]											=  (min_value >= expected_min) and (max_value <= expected_max)

											if not is_inside_interval:
																	lowercase       :    int											=  '''%s not found in %s''' % (safe_repr(lowerCAmelCase					), safe_repr(lowerCAmelCase					))
																	self.fail(self._formatMessage(lowerCAmelCase ,      lowerCAmelCase					)					)

					def        _lowerCAmelCase	(			self      :					Tuple					):
											self.config_tester.run_common_tests()

					def        _lowerCAmelCase	(			self      :					int					):
											lowercase       :    Optional[Any]											=  self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_rwkv_model(*lowerCAmelCase					)

					def        _lowerCAmelCase	(			self      :					Any					):
											lowercase       :    Optional[Any]											=  self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_causl_lm(*lowerCAmelCase					)

					def        _lowerCAmelCase	(			self      :					str					):
											lowercase       :    List[str]											=  self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_state_equivalency(*lowerCAmelCase					)

					def        _lowerCAmelCase	(			self      :					Tuple					):
											lowercase       :    Union[str, Any]											=  self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase       :    Union[str, Any]											=  model_class(config=lowerCAmelCase					)
																	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        _lowerCAmelCase	(			self      :					Optional[int]					):
											lowercase       :    Dict											=  self.model_tester.prepare_config_and_inputs_for_common()
											lowercase       :    Any											=  True

											lowercase       :    Tuple											=  getattr(self.model_tester ,      '''seq_length''' ,      lowerCAmelCase					)

											for model_class in self.all_model_classes:
																	lowercase       :    Dict											=  True
																	lowercase       :    int											=  False
																	lowercase       :    List[str]											=  True
																	lowercase       :    List[str]											=  model_class(lowerCAmelCase					)
																	model.to(lowerCAmelCase					)
																	model.eval()

																	lowercase       :    List[Any]											=  self._prepare_for_class(lowerCAmelCase ,      lowerCAmelCase					)
																	lowercase       :    int											=  inputs['''input_ids'''].shape[0]
																	with torch.no_grad():
																							lowercase       :    List[str]											=  model(**lowerCAmelCase					)
																	lowercase       :    List[Any]											=  outputs.attentions
																	self.assertEqual(len(lowerCAmelCase					) ,      self.model_tester.num_hidden_layers					)

																	# check that output_attentions also work using config
																	del inputs_dict["output_attentions"]
																	lowercase       :    Any											=  True
																	lowercase       :    List[str]											=  model_class(lowerCAmelCase					)
																	model.to(lowerCAmelCase					)
																	model.eval()

																	lowercase       :    str											=  self._prepare_for_class(lowerCAmelCase ,      lowerCAmelCase					)
																	lowercase       :    Any											=  inputs['''input_ids'''].shape[0]
																	with torch.no_grad():
																							lowercase       :    Union[str, Any]											=  model(**lowerCAmelCase					)
																	lowercase       :    Any											=  outputs.attentions
																	self.assertEqual(len(lowerCAmelCase					) ,      self.model_tester.num_hidden_layers					)

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

																	# Check attention is always last and order is fine
																	lowercase       :    Union[str, Any]											=  True
																	lowercase       :    Tuple											=  True
																	lowercase       :    str											=  model_class(lowerCAmelCase					)
																	model.to(lowerCAmelCase					)
																	model.eval()

																	lowercase       :    Tuple											=  self._prepare_for_class(lowerCAmelCase ,      lowerCAmelCase					)
																	lowercase       :    str											=  inputs['''input_ids'''].shape[0]
																	with torch.no_grad():
																							lowercase       :    Union[str, Any]											=  model(**lowerCAmelCase					)

																	lowercase       :    Dict											=  1
																	self.assertEqual(out_len + added_hidden_states ,      len(lowerCAmelCase					)					)

																	lowercase       :    Any											=  outputs.attentions

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




					@slow
					def        _lowerCAmelCase	(			self      :					Any					):
											for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase       :    List[str]											=  RwkvModel.from_pretrained(lowerCAmelCase					)
																	self.assertIsNotNone(lowerCAmelCase					)







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

					def        _lowerCAmelCase	(			self      :					Optional[int]					):
											lowercase       :    Optional[Any]											=  '''RWKV/rwkv-4-169m-pile'''
											lowercase       :    Optional[int]											=  AutoTokenizer.from_pretrained(self.model_id					)

					def        _lowerCAmelCase	(			self      :					int					):
											lowercase       :    Optional[int]											=  '''Hello my name is Jasmine and I am a newbie to the'''
											lowercase       :    Any											=  RwkvForCausalLM.from_pretrained(self.model_id					).to(lowerCAmelCase					)

											lowercase       :    Dict											=  self.tokenizer('''Hello my name is''' ,      return_tensors='''pt'''					).input_ids.to(lowerCAmelCase					)
											lowercase       :    str											=  model.generate(lowerCAmelCase ,      max_new_tokens=10					)
											lowercase       :    str											=  self.tokenizer.decode(output[0].tolist()					)

											self.assertEqual(lowerCAmelCase ,      lowerCAmelCase					)




					def        _lowerCAmelCase	(			self      :					Dict					):
											lowercase       :    List[str]											=  '''Hello my name is Jasmine and I am a newbie to the'''

											lowercase       :    Union[str, Any]											=  self.tokenizer('''Hello my name is''' ,      return_tensors='''pt'''					).input_ids.to(lowerCAmelCase					)
											lowercase       :    str											=  RwkvForCausalLM.from_pretrained(self.model_id ,      torch_dtype=torch.bfloataa					).to(lowerCAmelCase					)

											lowercase       :    List[Any]											=  model.generate(lowerCAmelCase ,      max_new_tokens=10					)
											lowercase       :    Optional[Any]											=  self.tokenizer.decode(output[0].tolist()					)

											self.assertEqual(lowerCAmelCase ,      lowerCAmelCase					)




from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import PNDMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class 							UpperCAmelCase	(   __lowerCamelCase					):
					a__: UNetaDModel
					a__: PNDMScheduler

					def __init__(			self      :					List[str] ,      lowerCAmelCase      :					UNetaDModel ,      lowerCAmelCase      :					PNDMScheduler					):
											super().__init__()

											lowercase       :    int											=  PNDMScheduler.from_config(scheduler.config					)

											self.register_modules(unet=lowerCAmelCase ,      scheduler=lowerCAmelCase					)




					@torch.no_grad()
					def __call__(			self      :					List[str] ,      lowerCAmelCase      :					int = 1 ,      lowerCAmelCase      :					int = 50 ,      lowerCAmelCase      :					Optional[Union[torch.Generator, List[torch.Generator]]] = None ,      lowerCAmelCase      :					Optional[str] = "pil" ,      lowerCAmelCase      :					bool = True ,      **lowerCAmelCase      :					int ,      ):
											lowercase       :    Tuple											=  randn_tensor(
											    (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size) ,      generator=lowerCAmelCase ,      device=self.device ,      )

											self.scheduler.set_timesteps(lowerCAmelCase					)
											for t in self.progress_bar(self.scheduler.timesteps					):
																	lowercase       :    Dict											=  self.unet(lowerCAmelCase ,      lowerCAmelCase					).sample

																	lowercase       :    Tuple											=  self.scheduler.step(lowerCAmelCase ,      lowerCAmelCase ,      lowerCAmelCase					).prev_sample

											lowercase       :    Optional[Any]											=  (image / 2 + 0.5).clamp(0 ,      1					)
											lowercase       :    Optional[int]											=  image.cpu().permute(0 ,      2 ,      3 ,      1					).numpy()
											if output_type == "pil":
																	lowercase       :    Union[str, Any]											=  self.numpy_to_pil(lowerCAmelCase					)

											if not return_dict:
																	return (image,)

											return ImagePipelineOutput(images=lowerCAmelCase					)


