






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




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




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




						def      snake_case__(    self		,	lowercase=False		,	lowercase=False		,	lowercase=False       )			->						List[Any]:
													_a  :    Optional[int]			=						ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size       )

													_a  :    List[Any]			=						None
													if self.use_input_mask:
																				_a  :    List[str]			=						random_attention_mask([self.batch_size, self.seq_length]       )

													_a  :    Dict			=						None
													if self.use_token_type_ids:
																				_a  :    Optional[int]			=						ids_tensor([self.batch_size, self.seq_length]		,	self.type_vocab_size       )

													_a  :    int			=						None
													if self.use_mc_token_ids:
																				_a  :    Union[str, Any]			=						ids_tensor([self.batch_size, self.num_choices]		,	self.seq_length       )

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

													_a  :    Union[str, Any]			=						self.get_config(
													    gradient_checkpointing=lowercase__		,	scale_attn_by_inverse_layer_idx=lowercase__		,	reorder_and_upcast_attn=lowercase__		,	)

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




						def      snake_case__(    self		,	lowercase=False		,	lowercase=False		,	lowercase=False       )			->						Optional[Any]:
													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=lowercase__		,	bos_token_id=self.bos_token_id		,	eos_token_id=self.eos_token_id		,	pad_token_id=self.pad_token_id		,	gradient_checkpointing=lowercase__		,	scale_attn_by_inverse_layer_idx=lowercase__		,	reorder_and_upcast_attn=lowercase__		,	)




						def      snake_case__(    self       )			->						Optional[int]:
													_a  :    Optional[int]			=						self.get_config()
													_a  :    int			=						300
													return config




						def      snake_case__(    self       )			->						Union[str, Any]:
													(
													    (
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					   
													)  :    int			=						self.prepare_config_and_inputs()

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

													_a  :    str			=						model(lowercase__       )

													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		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	*lowercase       )			->						List[str]:
													_a  :    List[str]			=						RwkvForCausalLM(lowercase__       )
													model.to(lowercase__       )
													model.eval()

													_a  :    List[str]			=						model(lowercase__		,	labels=lowercase__       )
													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		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	*lowercase       )			->						List[str]:
													_a  :    Union[str, Any]			=						RwkvModel(config=lowercase__       )
													model.to(lowercase__       )
													model.eval()

													_a  :    Tuple			=						model(lowercase__       )
													_a  :    Dict			=						outputs.last_hidden_state

													_a  :    List[str]			=						model(input_ids[:, :2]       )
													_a  :    int			=						outputs.last_hidden_state

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

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




						def      snake_case__(    self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	*lowercase		,	lowercase=False       )			->						Tuple:
													_a  :    Optional[Any]			=						RwkvForCausalLM(lowercase__       )
													model.to(lowercase__       )
													if gradient_checkpointing:
																				model.gradient_checkpointing_enable()

													_a  :    Optional[Any]			=						model(lowercase__		,	labels=lowercase__       )
													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       )			->						List[str]:
													_a  :    str			=						self.prepare_config_and_inputs()

													(
													    (
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					(
													    _a   
													)		,					   
													)  :    List[str]			=						config_and_inputs

													_a  :    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_							(				lowerCAmelCase__ ,       lowerCAmelCase__ ,       lowerCAmelCase__ ,       unittest.TestCase       ):
						lowercase									=  (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
						lowercase									=  (
						    {'feature-extraction': RwkvModel, 'text-generation': RwkvForCausalLM} if is_torch_available() else {}
						)
						# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
						lowercase									=  False
						lowercase									=  False
						lowercase									=  False
						lowercase									=  False
						lowercase									=  False  # Rwkv does not support head masking




						def      snake_case__(    self       )			->						str:
													_a  :    str			=						RwkvModelTester(self       )
													_a  :    Any			=						ConfigTester(
													    self		,	config_class=lowercase__		,	n_embd=37		,	common_properties=['''hidden_size''', '''num_hidden_layers''']       )




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

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

													_a		,					_a  :    Any			=						container

													_a  :    Any			=						(min_value >= expected_min) and (max_value <= expected_max)

													if not is_inside_interval:
																				_a  :    str			=						'''%s not found in %s''' % (safe_repr(lowercase__       ), safe_repr(lowercase__       ))
																				self.fail(self._formatMessage(lowercase__		,	lowercase__       )       )




						def      snake_case__(    self       )			->						Optional[Any]:
													self.config_tester.run_common_tests()




						def      snake_case__(    self       )			->						Union[str, Any]:
													_a  :    Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_rwkv_model(*lowercase__       )




						def      snake_case__(    self       )			->						Dict:
													_a  :    List[Any]			=						self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_causl_lm(*lowercase__       )




						def      snake_case__(    self       )			->						str:
													_a  :    Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_state_equivalency(*lowercase__       )




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

													for model_class in self.all_model_classes:
																				_a  :    Any			=						model_class(config=lowercase__       )
																				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       )			->						Tuple:
													_a		,					_a  :    str			=						self.model_tester.prepare_config_and_inputs_for_common()
													_a  :    Dict			=						True

													_a  :    List[Any]			=						getattr(self.model_tester		,	'''seq_length'''		,	lowercase__       )

													for model_class in self.all_model_classes:
																				_a  :    List[Any]			=						True
																				_a  :    List[str]			=						False
																				_a  :    Any			=						True
																				_a  :    Dict			=						model_class(lowercase__       )
																				model.to(lowercase__       )
																				model.eval()

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

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

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

																				self.assertListEqual(
																				    list(attentions[0].shape[-3:]       )		,	[batch_size, seq_len, config.hidden_size]		,	)
																				_a  :    Optional[Any]			=						len(lowercase__       )

																				# Check attention is always last and order is fine
																				_a  :    Optional[Any]			=						True
																				_a  :    Union[str, Any]			=						True
																				_a  :    Dict			=						model_class(lowercase__       )
																				model.to(lowercase__       )
																				model.eval()

																				_a  :    str			=						self._prepare_for_class(lowercase__		,	lowercase__       )
																				_a  :    Optional[int]			=						inputs['''input_ids'''].shape[0]
																				with torch.no_grad():
																											_a  :    Union[str, Any]			=						model(**lowercase__       )

																				_a  :    int			=						1
																				self.assertEqual(out_len + added_hidden_states		,	len(lowercase__       )       )

																				_a  :    List[Any]			=						outputs.attentions

																				self.assertEqual(len(lowercase__       )		,	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       )			->						List[str]:
													for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				_a  :    Any			=						RwkvModel.from_pretrained(lowercase__       )
																				self.assertIsNotNone(lowercase__       )

@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      snake_case__(    self       )			->						Dict:
													_a  :    Tuple			=						'''RWKV/rwkv-4-169m-pile'''
													_a  :    List[str]			=						AutoTokenizer.from_pretrained(self.model_id       )




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

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

													self.assertEqual(lowercase__		,	lowercase__       )







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

													_a  :    Any			=						self.tokenizer('''Hello my name is'''		,	return_tensors='''pt'''       ).input_ids.to(lowercase__       )
													_a  :    Any			=						RwkvForCausalLM.from_pretrained(self.model_id		,	torch_dtype=torch.bfloataa       ).to(lowercase__       )

													_a  :    Dict			=						model.generate(lowercase__		,	max_new_tokens=10       )
													_a  :    Optional[int]			=						self.tokenizer.decode(output[0].tolist()       )

													self.assertEqual(lowercase__		,	lowercase__       )







from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


snake_case_   :	Optional[Any]							      =logging.get_logger(__name__)

snake_case_   :	Any							      ={
    '''MBZUAI/swiftformer-xs''': '''https://huggingface.co/MBZUAI/swiftformer-xs/resolve/main/config.json''',
}





class        a__      (	lowerCAmelCase__		):
		UpperCAmelCase_							:       Union[str, Any] 					= 'swiftformer'






		def __init__(      self							, lowercase__=3							, lowercase__=[3, 3, 6, 4]							, lowercase__=[48, 56, 112, 220]							, lowercase__=4							, lowercase__=[True, True, True, True]							, lowercase__="gelu"							, lowercase__=3							, lowercase__=2							, lowercase__=1							, lowercase__=0.0							, lowercase__=True							, lowercase__=1e-5							, lowercase__=1e-5							, **lowercase__							, )				->				Optional[Any]:
				super().__init__(**lowercase__						)
				__A    =							num_channels
				__A    =							depths
				__A    =							embed_dims
				__A    =							mlp_ratio
				__A    =							downsamples
				__A    =							hidden_act
				__A    =							down_patch_size
				__A    =							down_stride
				__A    =							down_pad
				__A    =							drop_path_rate
				__A    =							use_layer_scale
				__A    =							layer_scale_init_value
				__A    =							batch_norm_eps





class        a__      (	lowerCAmelCase__		):
		UpperCAmelCase_							:       Union[str, Any] 					= version.parse('1.11'		)






		@property
		def 					_lowerCamelCase					(      self						)				->				Mapping[str, Mapping[int, str]]:
				return OrderedDict(
				    [
				        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
				    ]						)






		@property
		def 					_lowerCamelCase					(      self						)				->				float:
				return 1e-4

