

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_	     =     False





class      snake_case_       :







			'''simple docstring'''

			def __init__(					self			:    List[str]							,       _UpperCamelCase			:    Tuple							,       _UpperCamelCase			:    List[str]=1_4							,       _UpperCamelCase			:    List[Any]=7							,       _UpperCamelCase			:    Optional[int]=True							,       _UpperCamelCase			:    Tuple=False							,       _UpperCamelCase			:    Tuple=True							,       _UpperCamelCase			:    Union[str, Any]=True							,       _UpperCamelCase			:    List[Any]=True							,       _UpperCamelCase			:    Optional[Any]=9_9							,       _UpperCamelCase			:    Union[str, Any]=3_2							,       _UpperCamelCase			:    str=5							,       _UpperCamelCase			:    Optional[Any]=3_7							,       _UpperCamelCase			:    List[Any]="gelu"							,       _UpperCamelCase			:    int=0.1							,       _UpperCamelCase			:    int=0.1							,       _UpperCamelCase			:    Any=5_1_2							,       _UpperCamelCase			:    str=1_6							,       _UpperCamelCase			:    int=2							,       _UpperCamelCase			:    Optional[int]=3							,       _UpperCamelCase			:    Tuple=4							,       _UpperCamelCase			:    Tuple=None							,       )				->str:
									snake_case_       =     parent
									snake_case_       =     batch_size
									snake_case_       =     seq_length
									snake_case_       =     is_training
									snake_case_       =     use_token_type_ids
									snake_case_       =     use_input_mask
									snake_case_       =     use_labels
									snake_case_       =     use_mc_token_ids
									snake_case_       =     vocab_size
									snake_case_       =     hidden_size
									snake_case_       =     num_hidden_layers
									snake_case_       =     intermediate_size
									snake_case_       =     hidden_act
									snake_case_       =     hidden_dropout_prob
									snake_case_       =     attention_probs_dropout_prob
									snake_case_       =     max_position_embeddings
									snake_case_       =     type_vocab_size
									snake_case_       =     type_sequence_label_size
									snake_case_       =     num_labels
									snake_case_       =     num_choices
									snake_case_       =     scope
									snake_case_       =     vocab_size - 1
									snake_case_       =     vocab_size - 1
									snake_case_       =     vocab_size - 1







			def     snake_case__(					self			:    Tuple							)				->Tuple:
									return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''							)







			def     snake_case__(					self			:    List[str]							,       _UpperCamelCase			:    Any=False							,       _UpperCamelCase			:    str=False							,       _UpperCamelCase			:    int=False							)				->int:
									snake_case_       =     ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size							)

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

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

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

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

									snake_case_       =     self.get_config(
									    gradient_checkpointing=_UpperCamelCase							,       scale_attn_by_inverse_layer_idx=_UpperCamelCase							,       reorder_and_upcast_attn=_UpperCamelCase							,       )

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







			def     snake_case__(					self			:    Tuple							,       _UpperCamelCase			:    Optional[int]=False							,       _UpperCamelCase			:    Union[str, Any]=False							,       _UpperCamelCase			:    Optional[Any]=False							)				->Dict:
									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=_UpperCamelCase							,       bos_token_id=self.bos_token_id							,       eos_token_id=self.eos_token_id							,       pad_token_id=self.pad_token_id							,       gradient_checkpointing=_UpperCamelCase							,       scale_attn_by_inverse_layer_idx=_UpperCamelCase							,       reorder_and_upcast_attn=_UpperCamelCase							,       )







			def     snake_case__(					self			:    Optional[Any]							)				->Optional[int]:
									snake_case_       =     self.get_config()
									snake_case_       =     3_0_0
									return config







			def     snake_case__(					self			:    str							)				->List[Any]:
									(
									    (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   
									)       =     self.prepare_config_and_inputs()

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

									snake_case_       =     model(_UpperCamelCase							)

									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     snake_case__(					self			:    Union[str, Any]							,       _UpperCamelCase			:    List[str]							,       _UpperCamelCase			:    List[Any]							,       _UpperCamelCase			:    Optional[Any]							,       _UpperCamelCase			:    str							,       _UpperCamelCase			:    int							,       *_UpperCamelCase			:    Dict							)				->Any:
									snake_case_       =     RwkvForCausalLM(_UpperCamelCase							)
									model.to(_UpperCamelCase							)
									model.eval()

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







			def     snake_case__(					self			:    List[Any]							,       _UpperCamelCase			:    Optional[Any]							,       _UpperCamelCase			:    Any							,       _UpperCamelCase			:    Tuple							,       _UpperCamelCase			:    List[Any]							,       _UpperCamelCase			:    str							,       *_UpperCamelCase			:    Optional[Any]							)				->Union[str, Any]:
									snake_case_       =     RwkvModel(config=_UpperCamelCase							)
									model.to(_UpperCamelCase							)
									model.eval()

									snake_case_       =     model(_UpperCamelCase							)
									snake_case_       =     outputs.last_hidden_state

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

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

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







			def     snake_case__(					self			:    List[Any]							,       _UpperCamelCase			:    Union[str, Any]							,       _UpperCamelCase			:    Tuple							,       _UpperCamelCase			:    Optional[int]							,       _UpperCamelCase			:    int							,       _UpperCamelCase			:    Dict							,       *_UpperCamelCase			:    Dict							,       _UpperCamelCase			:    Union[str, Any]=False							)				->str:
									snake_case_       =     RwkvForCausalLM(_UpperCamelCase							)
									model.to(_UpperCamelCase							)
									if gradient_checkpointing:
															model.gradient_checkpointing_enable()

									snake_case_       =     model(_UpperCamelCase							,       labels=_UpperCamelCase							)
									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     snake_case__(					self			:    Dict							)				->List[Any]:
									snake_case_       =     self.prepare_config_and_inputs()

									(
									    (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   (
									    snake_case_
									),   
									)       =     config_and_inputs

									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 ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE					:       List[Any]						      =    (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
			SCREAMING_SNAKE_CASE					:       Tuple						      =    (
			    {"feature-extraction": RwkvModel, "text-generation": RwkvForCausalLM} if is_torch_available() else {}
			)
			# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
			SCREAMING_SNAKE_CASE					:       Dict						      =    False
			SCREAMING_SNAKE_CASE					:       Dict						      =    False
			SCREAMING_SNAKE_CASE					:       List[Any]						      =    False
			SCREAMING_SNAKE_CASE					:       int						      =    False
			SCREAMING_SNAKE_CASE					:       Optional[Any]						      =    False  # Rwkv does not support head masking







			def     snake_case__(					self			:    List[str]							)				->Optional[Any]:
									snake_case_       =     RwkvModelTester(self							)
									snake_case_       =     ConfigTester(
									    self							,       config_class=_UpperCamelCase							,       n_embd=3_7							,       common_properties=['''hidden_size''', '''num_hidden_layers''']							)







			def     snake_case__(					self			:    Tuple							,       _UpperCamelCase			:    int							,       _UpperCamelCase			:    Tuple							,       _UpperCamelCase			:    List[Any]=None							)				->Optional[Any]:
									if isinstance(_UpperCamelCase							,       torch.Tensor							):
															snake_case_,   snake_case_       =     member.max().item(), member.min().item()
									elif isinstance(_UpperCamelCase							,       _UpperCamelCase							) or isinstance(_UpperCamelCase							,       _UpperCamelCase							):
															snake_case_,   snake_case_       =     max(_UpperCamelCase							), min(_UpperCamelCase							)

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

									snake_case_,   snake_case_       =     container

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

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







			def     snake_case__(					self			:    Dict							)				->Tuple:
									self.config_tester.run_common_tests()







			def     snake_case__(					self			:    str							)				->str:
									snake_case_       =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_rwkv_model(*_UpperCamelCase							)







			def     snake_case__(					self			:    Dict							)				->str:
									snake_case_       =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_causl_lm(*_UpperCamelCase							)







			def     snake_case__(					self			:    Optional[Any]							)				->Any:
									snake_case_       =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_state_equivalency(*_UpperCamelCase							)







			def     snake_case__(					self			:    List[Any]							)				->List[str]:
									snake_case_,   snake_case_       =     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															snake_case_       =     model_class(config=_UpperCamelCase							)
															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     snake_case__(					self			:    List[str]							)				->Optional[int]:
									snake_case_,   snake_case_       =     self.model_tester.prepare_config_and_inputs_for_common()
									snake_case_       =     True

									snake_case_       =     getattr(self.model_tester							,       '''seq_length'''							,       _UpperCamelCase							)

									for model_class in self.all_model_classes:
															snake_case_       =     True
															snake_case_       =     False
															snake_case_       =     True
															snake_case_       =     model_class(_UpperCamelCase							)
															model.to(_UpperCamelCase							)
															model.eval()

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

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

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

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

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

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

															snake_case_       =     1
															self.assertEqual(out_len + added_hidden_states							,       len(_UpperCamelCase							)							)

															snake_case_       =     outputs.attentions

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







			@slow
			def     snake_case__(					self			:    Any							)				->Optional[int]:
									for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															snake_case_       =     RwkvModel.from_pretrained(_UpperCamelCase							)
															self.assertIsNotNone(_UpperCamelCase							)






@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 ):







			'''simple docstring'''

			def     snake_case__(					self			:    Dict							)				->List[Any]:
									snake_case_       =     '''RWKV/rwkv-4-169m-pile'''
									snake_case_       =     AutoTokenizer.from_pretrained(self.model_id							)







			def     snake_case__(					self			:    Optional[Any]							)				->Optional[Any]:
									snake_case_       =     '''Hello my name is Jasmine and I am a newbie to the'''
									snake_case_       =     RwkvForCausalLM.from_pretrained(self.model_id							).to(_UpperCamelCase							)

									snake_case_       =     self.tokenizer('''Hello my name is'''							,       return_tensors='''pt'''							).input_ids.to(_UpperCamelCase							)
									snake_case_       =     model.generate(_UpperCamelCase							,       max_new_tokens=1_0							)
									snake_case_       =     self.tokenizer.decode(output[0].tolist()							)

									self.assertEqual(_UpperCamelCase							,       _UpperCamelCase							)







			def     snake_case__(					self			:    List[Any]							)				->Any:
									snake_case_       =     '''Hello my name is Jasmine and I am a newbie to the'''

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

									snake_case_       =     model.generate(_UpperCamelCase							,       max_new_tokens=1_0							)
									snake_case_       =     self.tokenizer.decode(output[0].tolist()							)

									self.assertEqual(_UpperCamelCase							,       _UpperCamelCase							)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__	):
						snake_case_       =     5381
						for x in s:
												snake_case_       =     ((hash_value << 5) + hash_value) + ord(SCREAMING_SNAKE_CASE__	)
						return hash_value & 0xFF_FF_FF_FF