

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


class 			_SCREAMING_SNAKE_CASE      :







						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     ,						)->							str:
													lowerCamelCase_								=parent
													lowerCamelCase_								=batch_size
													lowerCamelCase_								=seq_length
													lowerCamelCase_								=is_training
													lowerCamelCase_								=use_token_type_ids
													lowerCamelCase_								=use_input_mask
													lowerCamelCase_								=use_labels
													lowerCamelCase_								=use_mc_token_ids
													lowerCamelCase_								=vocab_size
													lowerCamelCase_								=hidden_size
													lowerCamelCase_								=num_hidden_layers
													lowerCamelCase_								=intermediate_size
													lowerCamelCase_								=hidden_act
													lowerCamelCase_								=hidden_dropout_prob
													lowerCamelCase_								=attention_probs_dropout_prob
													lowerCamelCase_								=max_position_embeddings
													lowerCamelCase_								=type_vocab_size
													lowerCamelCase_								=type_sequence_label_size
													lowerCamelCase_								=num_labels
													lowerCamelCase_								=num_choices
													lowerCamelCase_								=scope
													lowerCamelCase_								=vocab_size - 1
													lowerCamelCase_								=vocab_size - 1
													lowerCamelCase_								=vocab_size - 1







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







						def        _snake_case				(			self     ,						_SCREAMING_SNAKE_CASE=False     ,						_SCREAMING_SNAKE_CASE=False     ,						_SCREAMING_SNAKE_CASE=False							)->							Any:
													lowerCamelCase_								=ids_tensor([self.batch_size, self.seq_length]     ,						self.vocab_size							)

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

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

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

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

													lowerCamelCase_								=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     ,						_SCREAMING_SNAKE_CASE=False     ,						_SCREAMING_SNAKE_CASE=False     ,						_SCREAMING_SNAKE_CASE=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							)->							int:
													lowerCamelCase_								=self.get_config()
													lowerCamelCase_								=300
													return config







						def        _snake_case				(			self							)->							Union[str, Any]:
													(
													    (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   			
													)								=self.prepare_config_and_inputs()

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

													lowerCamelCase_								=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     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						*_SCREAMING_SNAKE_CASE							)->							Dict:
													lowerCamelCase_								=RwkvForCausalLM(lowercase_							)
													model.to(lowercase_							)
													model.eval()

													lowerCamelCase_								=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     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						*_SCREAMING_SNAKE_CASE							)->							int:
													lowerCamelCase_								=RwkvModel(config=lowercase_							)
													model.to(lowercase_							)
													model.eval()

													lowerCamelCase_								=model(lowercase_							)
													lowerCamelCase_								=outputs.last_hidden_state

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

													# Using the state computed on the first inputs, we will get the same output
													lowerCamelCase_								=model(input_ids[:, 2:]     ,						state=outputs.state							)
													lowerCamelCase_								=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     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						*_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE=False							)->							List[str]:
													lowerCamelCase_								=RwkvForCausalLM(lowercase_							)
													model.to(lowercase_							)
													if gradient_checkpointing:
																				model.gradient_checkpointing_enable()

													lowerCamelCase_								=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							)->							Any:
													lowerCamelCase_								=self.prepare_config_and_inputs()

													(
													    (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   (
													    lowerCamelCase_			
													)	,   			
													)								=config_and_inputs

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







						def        _snake_case				(			self							)->							Any:
													lowerCamelCase_								=RwkvModelTester(self							)
													lowerCamelCase_								=ConfigTester(
													    self     ,						config_class=lowercase_     ,						n_embd=37     ,						common_properties=["""hidden_size""", """num_hidden_layers"""]							)







						def        _snake_case				(			self     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE     ,						_SCREAMING_SNAKE_CASE=None							)->							Union[str, Any]:
													if isinstance(lowercase_     ,						torch.Tensor							):
																				lowerCamelCase_	,   lowerCamelCase_								=member.max().item(), member.min().item()
													elif isinstance(lowercase_     ,						lowercase_							) or isinstance(lowercase_     ,						lowercase_							):
																				lowerCamelCase_	,   lowerCamelCase_								=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"""							)

													lowerCamelCase_	,   lowerCamelCase_								=container

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

													if not is_inside_interval:
																				lowerCamelCase_								="""%s not found in %s""" % (safe_repr(lowercase_							), safe_repr(lowercase_							))
																				self.fail(self._formatMessage(lowercase_     ,						lowercase_							)							)







						def        _snake_case				(			self							)->							Any:
													self.config_tester.run_common_tests()







						def        _snake_case				(			self							)->							Optional[int]:
													lowerCamelCase_								=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_rwkv_model(*lowercase_							)







						def        _snake_case				(			self							)->							int:
													lowerCamelCase_								=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_causl_lm(*lowercase_							)







						def        _snake_case				(			self							)->							Any:
													lowerCamelCase_								=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_state_equivalency(*lowercase_							)







						def        _snake_case				(			self							)->							Optional[Any]:
													lowerCamelCase_	,   lowerCamelCase_								=self.model_tester.prepare_config_and_inputs_for_common()

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

													lowerCamelCase_								=getattr(self.model_tester     ,						"""seq_length"""     ,						lowercase_							)

													for model_class in self.all_model_classes:
																				lowerCamelCase_								=True
																				lowerCamelCase_								=False
																				lowerCamelCase_								=True
																				lowerCamelCase_								=model_class(lowercase_							)
																				model.to(lowercase_							)
																				model.eval()

																				lowerCamelCase_								=self._prepare_for_class(lowercase_     ,						lowercase_							)
																				lowerCamelCase_								=inputs["""input_ids"""].shape[0]
																				with torch.no_grad():
																											lowerCamelCase_								=model(**lowercase_							)
																				lowerCamelCase_								=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"]
																				lowerCamelCase_								=True
																				lowerCamelCase_								=model_class(lowercase_							)
																				model.to(lowercase_							)
																				model.eval()

																				lowerCamelCase_								=self._prepare_for_class(lowercase_     ,						lowercase_							)
																				lowerCamelCase_								=inputs["""input_ids"""].shape[0]
																				with torch.no_grad():
																											lowerCamelCase_								=model(**lowercase_							)
																				lowerCamelCase_								=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]     ,						)
																				lowerCamelCase_								=len(lowercase_							)

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

																				lowerCamelCase_								=self._prepare_for_class(lowercase_     ,						lowercase_							)
																				lowerCamelCase_								=inputs["""input_ids"""].shape[0]
																				with torch.no_grad():
																											lowerCamelCase_								=model(**lowercase_							)

																				lowerCamelCase_								=1
																				self.assertEqual(out_len + added_hidden_states     ,						len(lowercase_							)							)

																				lowerCamelCase_								=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							)->							str:
													for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				lowerCamelCase_								=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 			_SCREAMING_SNAKE_CASE      (  unittest.TestCase):







						def        _snake_case				(			self							)->							Optional[int]:
													lowerCamelCase_								="""RWKV/rwkv-4-169m-pile"""
													lowerCamelCase_								=AutoTokenizer.from_pretrained(self.model_id							)







						def        _snake_case				(			self							)->							List[str]:
													lowerCamelCase_								="""Hello my name is Jasmine and I am a newbie to the"""
													lowerCamelCase_								=RwkvForCausalLM.from_pretrained(self.model_id							).to(lowercase_							)

													lowerCamelCase_								=self.tokenizer("""Hello my name is"""     ,						return_tensors="""pt"""							).input_ids.to(lowercase_							)
													lowerCamelCase_								=model.generate(lowercase_     ,						max_new_tokens=10							)
													lowerCamelCase_								=self.tokenizer.decode(output[0].tolist()							)

													self.assertEqual(lowercase_     ,						lowercase_							)







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

													lowerCamelCase_								=self.tokenizer("""Hello my name is"""     ,						return_tensors="""pt"""							).input_ids.to(lowercase_							)
													lowerCamelCase_								=RwkvForCausalLM.from_pretrained(self.model_id     ,						torch_dtype=torch.bfloataa							).to(lowercase_							)

													lowerCamelCase_								=model.generate(lowercase_     ,						max_new_tokens=10							)
													lowerCamelCase_								=self.tokenizer.decode(output[0].tolist()							)

													self.assertEqual(lowercase_     ,						lowercase_							)



'''simple docstring'''

lowerCAmelCase_     =      "\n# Installation de Transformers\n! pip install transformers datasets\n# Pour installer à partir du code source au lieu de la dernière version, commentez la commande ci-dessus et décommentez la suivante.\n# ! pip install git+https://github.com/huggingface/transformers.git\n"

lowerCAmelCase_     =      [{"type": "code", "content": INSTALL_CONTENT}]
lowerCAmelCase_     =      {
    "{processor_class}": "FakeProcessorClass",
    "{model_class}": "FakeModelClass",
    "{object_class}": "FakeObjectClass",
}





