




"""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:
							__A							=     False
class      lowerCamelCase__      :


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



					"""simple docstring"""







					snake_case						:  Dict        =			parent
					snake_case						:  List[str]        =			batch_size
					snake_case						:  int        =			seq_length
					snake_case						:  str        =			is_training
					snake_case						:  int        =			use_token_type_ids
					snake_case						:  List[str]        =			use_input_mask
					snake_case						:  Dict        =			use_labels
					snake_case						:  str        =			use_mc_token_ids
					snake_case						:  int        =			vocab_size
					snake_case						:  List[str]        =			hidden_size
					snake_case						:  Optional[int]        =			num_hidden_layers
					snake_case						:  List[str]        =			intermediate_size
					snake_case						:  Tuple        =			hidden_act
					snake_case						:  Optional[Any]        =			hidden_dropout_prob
					snake_case						:  List[str]        =			attention_probs_dropout_prob
					snake_case						:  List[Any]        =			max_position_embeddings
					snake_case						:  str        =			type_vocab_size
					snake_case						:  Tuple        =			type_sequence_label_size
					snake_case						:  List[str]        =			num_labels
					snake_case						:  Optional[int]        =			num_choices
					snake_case						:  Optional[Any]        =			scope
					snake_case						:  Optional[Any]        =			vocab_size - 1
					snake_case						:  str        =			vocab_size - 1
					snake_case						:  Tuple        =			vocab_size - 1


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE=False			,	SCREAMING_SNAKE_CASE=False			,	SCREAMING_SNAKE_CASE=False							):



					"""simple docstring"""







					snake_case						:  Optional[Any]        =			ids_tensor([self.batch_size, self.seq_length]			,	self.vocab_size							)

					snake_case						:  Tuple        =			None
					if self.use_input_mask:
							snake_case						:  Dict        =			random_attention_mask([self.batch_size, self.seq_length]							)

					snake_case						:  Tuple        =			None
					if self.use_token_type_ids:
							snake_case						:  Optional[int]        =			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[int]        =			ids_tensor([self.batch_size, self.num_choices]			,	self.seq_length							)

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

					snake_case						:  Union[str, Any]        =			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							):



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



					"""simple docstring"""







					snake_case						:  Optional[int]        =			self.get_config()
					snake_case						:  Tuple        =			300
					return config


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					(
					    snake_case			
					)						:  Any        =			self.prepare_config_and_inputs()

					snake_case						:  Optional[Any]        =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
					snake_case						:  Union[str, 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			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	*SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Dict        =			True
					snake_case						:  Optional[int]        =			RwkvModel(config=SCREAMING_SNAKE_CASE							)
					model.to(SCREAMING_SNAKE_CASE							)
					model.eval()

					snake_case						:  List[str]        =			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							):



					"""simple docstring"""







					snake_case						:  Optional[int]        =			RwkvForCausalLM(SCREAMING_SNAKE_CASE							)
					model.to(SCREAMING_SNAKE_CASE							)
					model.eval()

					snake_case						:  Union[str, 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)							)


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	*SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Tuple        =			RwkvModel(config=SCREAMING_SNAKE_CASE							)
					model.to(SCREAMING_SNAKE_CASE							)
					model.eval()

					snake_case						:  List[str]        =			model(SCREAMING_SNAKE_CASE							)
					snake_case						:  Any        =			outputs.last_hidden_state

					snake_case						:  Union[str, Any]        =			model(input_ids[:, :2]							)
					snake_case						:  Dict        =			outputs.last_hidden_state

					# Using the state computed on the first inputs, we will get the same output
					snake_case						:  Tuple        =			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							):



					"""simple docstring"""







					snake_case						:  Optional[Any]        =			RwkvForCausalLM(SCREAMING_SNAKE_CASE							)
					model.to(SCREAMING_SNAKE_CASE							)
					if gradient_checkpointing:
							model.gradient_checkpointing_enable()

					snake_case						:  Union[str, 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							):



					"""simple docstring"""







					snake_case						:  List[Any]        =			self.prepare_config_and_inputs()

					(
					    snake_case			
					)						:  List[str]        =			config_and_inputs

					snake_case						:  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      lowerCamelCase__      (					lowerCamelCase_			,					lowerCamelCase_			,					lowerCamelCase_			,					unittest.TestCase     ):
			a__	: Tuple        =       (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
			a__	: Optional[int]        =       (
			    {"""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__	: List[Any]        =       False
			a__	: Union[str, Any]        =       False
			a__	: Optional[Any]        =       False
			a__	: Union[str, Any]        =       False  # Rwkv does not support head masking


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  List[Any]        =			RwkvModelTester(self							)
					snake_case						:  str        =			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							):



					"""simple docstring"""







					if isinstance(SCREAMING_SNAKE_CASE			,	torch.Tensor							):
							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						:  List[str]        =			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						:  str        =			container

					snake_case						:  int        =			(min_value >= expected_min) and (max_value <= expected_max)

					if not is_inside_interval:
							snake_case						:  Any        =			"%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							):



					"""simple docstring"""







					self.config_tester.run_common_tests()


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Tuple        =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_rwkv_model(*SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  List[Any]        =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_causl_lm(*SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Tuple        =			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_state_equivalency(*SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  List[Any]        =			self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							snake_case						:  List[str]        =			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							):



					"""simple docstring"""







					snake_case						:  Tuple        =			self.model_tester.prepare_config_and_inputs_for_common()
					snake_case						:  int        =			True

					snake_case						:  Dict        =			getattr(self.model_tester			,	"seq_length"			,	SCREAMING_SNAKE_CASE							)

					for model_class in self.all_model_classes:
							snake_case						:  Dict        =			True
							snake_case						:  List[str]        =			False
							snake_case						:  Union[str, Any]        =			True
							snake_case						:  Any        =			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[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						:  List[Any]        =			True
							snake_case						:  List[Any]        =			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						:  Union[str, Any]        =			inputs["input_ids"].shape[0]
							with torch.no_grad():
									snake_case						:  Optional[int]        =			model(**SCREAMING_SNAKE_CASE							)
							snake_case						:  List[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						:  str        =			len(SCREAMING_SNAKE_CASE							)

							# Check attention is always last and order is fine
							snake_case						:  List[Any]        =			True
							snake_case						:  Dict        =			True
							snake_case						:  str        =			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						:  Optional[Any]        =			inputs["input_ids"].shape[0]
							with torch.no_grad():
									snake_case						:  Optional[Any]        =			model(**SCREAMING_SNAKE_CASE							)

							snake_case						:  int        =			1
							self.assertEqual(out_len + added_hidden_states			,	len(SCREAMING_SNAKE_CASE							)							)

							snake_case						:  List[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							):



					"""simple docstring"""







					for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							snake_case						:  Optional[Any]        =			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      lowerCamelCase__      (					unittest.TestCase     ):


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Dict        =			"RWKV/rwkv-4-169m-pile"
					snake_case						:  Optional[int]        =			AutoTokenizer.from_pretrained(self.model_id							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Tuple        =			"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						:  Optional[Any]        =			self.tokenizer("Hello my name is"			,	return_tensors="pt"							).input_ids.to(SCREAMING_SNAKE_CASE							)
					snake_case						:  Union[str, Any]        =			model.generate(SCREAMING_SNAKE_CASE			,	max_new_tokens=10							)
					snake_case						:  int        =			self.tokenizer.decode(output[0].tolist()							)

					self.assertEqual(SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							)



			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  List[Any]        =			"Hello my name is Jasmine and I am a newbie to the"

					snake_case						:  Any        =			self.tokenizer("Hello my name is"			,	return_tensors="pt"							).input_ids.to(SCREAMING_SNAKE_CASE							)
					snake_case						:  Optional[int]        =			RwkvForCausalLM.from_pretrained(self.model_id			,	torch_dtype=torch.bfloataa							).to(SCREAMING_SNAKE_CASE							)

					snake_case						:  str        =			model.generate(SCREAMING_SNAKE_CASE			,	max_new_tokens=10							)
					snake_case						:  Any        =			self.tokenizer.decode(output[0].tolist()							)

					self.assertEqual(SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							)







"""simple docstring"""

def 					UpperCamelCase__       (						lowercase__					:       set    ,       lowercase__					:       set				):
		assert isinstance(lowercase__    ,       lowercase__				), F'''The input value of [set_a={set_a}] is not a set'''
		assert isinstance(lowercase__    ,       lowercase__				), F'''The input value of [set_b={set_b}] is not a set'''

		return {a + b for a in set_a for b in set_b}


if __name__ == "__main__":
							from doctest import testmod

							testmod()

