






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:
							lowercase     :							List[Any]	     =			False




class 						UpperCAmelCase_  :


						'''simple docstring'''




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






						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								return RwkvConfig.from_pretrained("sgugger/rwkv-4-pile-7b" )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE=False						, _SCREAMING_SNAKE_CASE=False						, _SCREAMING_SNAKE_CASE=False )       -> Optional[int]:
								snake_case_ :							Any         =				ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size )

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

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

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

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

								snake_case_ :							List[str]         =				self.get_config(
								    gradient_checkpointing=_SCREAMING_SNAKE_CASE						, scale_attn_by_inverse_layer_idx=_SCREAMING_SNAKE_CASE						, reorder_and_upcast_attn=_SCREAMING_SNAKE_CASE						, )

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE=False						, _SCREAMING_SNAKE_CASE=False						, _SCREAMING_SNAKE_CASE=False )       -> Optional[int]:
								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=_SCREAMING_SNAKE_CASE						, bos_token_id=self.bos_token_id						, eos_token_id=self.eos_token_id						, pad_token_id=self.pad_token_id						, gradient_checkpointing=_SCREAMING_SNAKE_CASE						, scale_attn_by_inverse_layer_idx=_SCREAMING_SNAKE_CASE						, reorder_and_upcast_attn=_SCREAMING_SNAKE_CASE						, )






						def 					_lowerCAmelCase				(		self )       -> Any:
								snake_case_ :							Optional[int]         =				self.get_config()
								snake_case_ :							int         =				300
								return config






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								(
								    (
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				
								) :							Optional[Any]         =				self.prepare_config_and_inputs()

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

								snake_case_ :							Any         =				model(_SCREAMING_SNAKE_CASE )

								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						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, *_SCREAMING_SNAKE_CASE )       -> Union[str, Any]:
								snake_case_ :							List[str]         =				RwkvForCausalLM(_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()

								snake_case_ :							Optional[int]         =				model(_SCREAMING_SNAKE_CASE						, labels=_SCREAMING_SNAKE_CASE )
								self.parent.assertEqual(result.loss.shape						, () )
								self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size) )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, *_SCREAMING_SNAKE_CASE )       -> List[Any]:
								snake_case_ :							int         =				RwkvModel(config=_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()

								snake_case_ :							Optional[Any]         =				model(_SCREAMING_SNAKE_CASE )
								snake_case_ :							Dict         =				outputs.last_hidden_state

								snake_case_ :							Union[str, Any]         =				model(input_ids[:, :2] )
								snake_case_ :							Optional[int]         =				outputs.last_hidden_state

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

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, *_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE=False )       -> Optional[Any]:
								snake_case_ :							Optional[int]         =				RwkvForCausalLM(_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								if gradient_checkpointing:
										model.gradient_checkpointing_enable()

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

								(
								    (
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				
								) :							List[Any]         =				config_and_inputs

								snake_case_ :							Union[str, Any]         =				{"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_  (		SCREAMING_SNAKE_CASE__    ,     SCREAMING_SNAKE_CASE__    ,     SCREAMING_SNAKE_CASE__    ,     unittest.TestCase    ):


						'''simple docstring'''




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






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_ :							Optional[Any]         =				RwkvModelTester(self )
								snake_case_ :							Optional[Any]         =				ConfigTester(
								    self						, config_class=_SCREAMING_SNAKE_CASE						, n_embd=37						, common_properties=["hidden_size", "num_hidden_layers"] )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE=None )       -> Optional[int]:
								if isinstance(_SCREAMING_SNAKE_CASE						, torch.Tensor ):
										snake_case_       ,				snake_case_ :							List[str]         =				member.max().item(), member.min().item()
								elif isinstance(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ) or isinstance(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ):
										snake_case_       ,				snake_case_ :							Dict         =				max(_SCREAMING_SNAKE_CASE ), min(_SCREAMING_SNAKE_CASE )

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

								snake_case_       ,				snake_case_ :							Optional[Any]         =				container

								snake_case_ :							Optional[Any]         =				(min_value >= expected_min) and (max_value <= expected_max)

								if not is_inside_interval:
										snake_case_ :							str         =				"%s not found in %s" % (safe_repr(_SCREAMING_SNAKE_CASE ), safe_repr(_SCREAMING_SNAKE_CASE ))
										self.fail(self._formatMessage(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ) )






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								self.config_tester.run_common_tests()






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_ :							List[Any]         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_rwkv_model(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> int:
								snake_case_ :							str         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_causl_lm(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Tuple:
								snake_case_ :							Optional[Any]         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_state_equivalency(*_SCREAMING_SNAKE_CASE )






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

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

								snake_case_ :							Tuple         =				getattr(self.model_tester						, "seq_length"						, _SCREAMING_SNAKE_CASE )

								for model_class in self.all_model_classes:
										snake_case_ :							str         =				True
										snake_case_ :							List[Any]         =				False
										snake_case_ :							Tuple         =				True
										snake_case_ :							Optional[int]         =				model_class(_SCREAMING_SNAKE_CASE )
										model.to(_SCREAMING_SNAKE_CASE )
										model.eval()

										snake_case_ :							Dict         =				self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
										snake_case_ :							Optional[Any]         =				inputs["input_ids"].shape[0]
										with torch.no_grad():
												snake_case_ :							Union[str, Any]         =				model(**_SCREAMING_SNAKE_CASE )
										snake_case_ :							Tuple         =				outputs.attentions
										self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, self.model_tester.num_hidden_layers )

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

										snake_case_ :							Any         =				self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
										snake_case_ :							int         =				inputs["input_ids"].shape[0]
										with torch.no_grad():
												snake_case_ :							Optional[int]         =				model(**_SCREAMING_SNAKE_CASE )
										snake_case_ :							Optional[Any]         =				outputs.attentions
										self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, self.model_tester.num_hidden_layers )

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

										# Check attention is always last and order is fine
										snake_case_ :							Optional[Any]         =				True
										snake_case_ :							Optional[int]         =				True
										snake_case_ :							int         =				model_class(_SCREAMING_SNAKE_CASE )
										model.to(_SCREAMING_SNAKE_CASE )
										model.eval()

										snake_case_ :							List[str]         =				self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
										snake_case_ :							Any         =				inputs["input_ids"].shape[0]
										with torch.no_grad():
												snake_case_ :							Optional[int]         =				model(**_SCREAMING_SNAKE_CASE )

										snake_case_ :							Optional[int]         =				1
										self.assertEqual(out_len + added_hidden_states						, len(_SCREAMING_SNAKE_CASE ) )

										snake_case_ :							Union[str, Any]         =				outputs.attentions

										self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, 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 )       -> Dict:
								for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										snake_case_ :							Optional[int]         =				RwkvModel.from_pretrained(_SCREAMING_SNAKE_CASE )
										self.assertIsNotNone(_SCREAMING_SNAKE_CASE )




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


						'''simple docstring'''




						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_ :							List[Any]         =				"RWKV/rwkv-4-169m-pile"
								snake_case_ :							Tuple         =				AutoTokenizer.from_pretrained(self.model_id )






						def 					_lowerCAmelCase				(		self )       -> Tuple:
								snake_case_ :							Any         =				"Hello my name is Jasmine and I am a newbie to the"
								snake_case_ :							int         =				RwkvForCausalLM.from_pretrained(self.model_id ).to(_SCREAMING_SNAKE_CASE )

								snake_case_ :							Tuple         =				self.tokenizer("Hello my name is"						, return_tensors="pt" ).input_ids.to(_SCREAMING_SNAKE_CASE )
								snake_case_ :							Optional[int]         =				model.generate(_SCREAMING_SNAKE_CASE						, max_new_tokens=10 )
								snake_case_ :							Optional[int]         =				self.tokenizer.decode(output[0].tolist() )

								self.assertEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> str:
								snake_case_ :							Union[str, Any]         =				"Hello my name is Jasmine and I am a newbie to the"

								snake_case_ :							str         =				self.tokenizer("Hello my name is"						, return_tensors="pt" ).input_ids.to(_SCREAMING_SNAKE_CASE )
								snake_case_ :							List[str]         =				RwkvForCausalLM.from_pretrained(self.model_id						, torch_dtype=torch.bfloataa ).to(_SCREAMING_SNAKE_CASE )

								snake_case_ :							List[str]         =				model.generate(_SCREAMING_SNAKE_CASE						, max_new_tokens=10 )
								snake_case_ :							List[str]         =				self.tokenizer.decode(output[0].tolist() )

								self.assertEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )








def    lowerCAmelCase__		(			_a				:	int			):
		if not isinstance(_a     ,					_a			):
				snake_case_ :							Optional[int]         =				F'''Input value of [number={number}] must be an integer'''
				raise TypeError(_a			)
		if number < 1:
				raise ValueError("Input must be a positive integer"			)
		return number * (2 * number - 1)


if __name__ == "__main__":
							import doctest

							doctest.testmod()
