






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:
  SCREAMING_SNAKE_CASE								=  False
class 							A_      :






 '''simple docstring'''




 def __init__(		self							,					_A							,					_A=14							,					_A=7							,					_A=True							,					_A=False							,					_A=True							,					_A=True							,					_A=True							,					_A=99							,					_A=32							,					_A=5							,					_A=37							,					_A="gelu"							,					_A=0.1							,					_A=0.1							,					_A=512							,					_A=16							,					_A=2							,					_A=3							,					_A=4							,					_A=None							,					)	->						str:

     """simple docstring"""



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


 def 	snake_case__			(		self)	->						int:

     """simple docstring"""



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


 def 	snake_case__			(		self							,					_A=False							,					_A=False							,					_A=False)	->						List[Any]:

     """simple docstring"""



     _UpperCAmelCase			: Tuple     					=			ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size)

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

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

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

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

     _UpperCAmelCase			: List[Any]     					=			self.get_config(
         gradient_checkpointing=_A							,					scale_attn_by_inverse_layer_idx=_A							,					reorder_and_upcast_attn=_A							,					)

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


 def 	snake_case__			(		self							,					_A=False							,					_A=False							,					_A=False)	->						Union[str, Any]:

     """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=_A							,					bos_token_id=self.bos_token_id							,					eos_token_id=self.eos_token_id							,					pad_token_id=self.pad_token_id							,					gradient_checkpointing=_A							,					scale_attn_by_inverse_layer_idx=_A							,					reorder_and_upcast_attn=_A							,					)


 def 	snake_case__			(		self)	->						Any:

     """simple docstring"""



     _UpperCAmelCase			: Dict     					=			self.get_config()
     _UpperCAmelCase			: Optional[Any]     					=			300
     return config


 def 	snake_case__			(		self)	->						Dict:

     """simple docstring"""



     (
         (
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	
     )			: Dict     					=			self.prepare_config_and_inputs()

     _UpperCAmelCase			: List[Any]     					=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
     _UpperCAmelCase			: Optional[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 	snake_case__			(		self							,					_A							,					_A							,					_A							,					_A							,					_A							,					*_A)	->						List[Any]:

     """simple docstring"""



     _UpperCAmelCase			: Optional[Any]     					=			True
     _UpperCAmelCase			: Optional[Any]     					=			RwkvModel(config=_A)
     model.to(_A)
     model.eval()

     _UpperCAmelCase			: Optional[Any]     					=			model(_A)

     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							,					_A							,					_A							,					_A							,					_A							,					_A							,					*_A)	->						Optional[int]:

     """simple docstring"""



     _UpperCAmelCase			: Dict     					=			RwkvForCausalLM(_A)
     model.to(_A)
     model.eval()

     _UpperCAmelCase			: List[str]     					=			model(_A							,					labels=_A)
     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							,					_A							,					_A							,					_A							,					_A							,					_A							,					*_A)	->						Tuple:

     """simple docstring"""



     _UpperCAmelCase			: Any     					=			RwkvModel(config=_A)
     model.to(_A)
     model.eval()

     _UpperCAmelCase			: Optional[int]     					=			model(_A)
     _UpperCAmelCase			: int     					=			outputs.last_hidden_state

     _UpperCAmelCase			: Optional[Any]     					=			model(input_ids[:, :2])
     _UpperCAmelCase			: Tuple     					=			outputs.last_hidden_state

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

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


 def 	snake_case__			(		self							,					_A							,					_A							,					_A							,					_A							,					_A							,					*_A							,					_A=False)	->						Dict:

     """simple docstring"""



     _UpperCAmelCase			: List[Any]     					=			RwkvForCausalLM(_A)
     model.to(_A)
     if gradient_checkpointing:
         model.gradient_checkpointing_enable()

     _UpperCAmelCase			: str     					=			model(_A							,					labels=_A)
     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)	->						Optional[Any]:

     """simple docstring"""



     _UpperCAmelCase			: Union[str, Any]     					=			self.prepare_config_and_inputs()

     (
         (
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	(
         _UpperCAmelCase
     )					,	
     )			: Tuple     					=			config_and_inputs

     _UpperCAmelCase			: 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 							A_      (  __lowercase	,				__lowercase	,				__lowercase	,				unittest.TestCase     ):






 '''simple docstring'''




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


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

     """simple docstring"""



     _UpperCAmelCase			: int     					=			RwkvModelTester(self)
     _UpperCAmelCase			: Optional[int]     					=			ConfigTester(
         self							,					config_class=_A							,					n_embd=37							,					common_properties=['''hidden_size''', '''num_hidden_layers'''])


 def 	snake_case__			(		self							,					_A							,					_A							,					_A=None)	->						List[str]:

     """simple docstring"""



     if isinstance(_A							,					torch.Tensor):
         _UpperCAmelCase					,	_UpperCAmelCase			: List[Any]     					=			member.max().item(), member.min().item()
     elif isinstance(_A							,					_A) or isinstance(_A							,					_A):
         _UpperCAmelCase					,	_UpperCAmelCase			: int     					=			max(_A), min(_A)

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

     _UpperCAmelCase					,	_UpperCAmelCase			: int     					=			container

     _UpperCAmelCase			: List[str]     					=			(min_value >= expected_min) and (max_value <= expected_max)

     if not is_inside_interval:
         _UpperCAmelCase			: Optional[Any]     					=			'''%s not found in %s''' % (safe_repr(_A), safe_repr(_A))
         self.fail(self._formatMessage(_A							,					_A))


 def 	snake_case__			(		self)	->						Union[str, Any]:

     """simple docstring"""



     self.config_tester.run_common_tests()


 def 	snake_case__			(		self)	->						List[str]:

     """simple docstring"""



     _UpperCAmelCase			: Union[str, Any]     					=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_rwkv_model(*_A)


 def 	snake_case__			(		self)	->						List[Any]:

     """simple docstring"""



     _UpperCAmelCase			: Optional[int]     					=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_causl_lm(*_A)


 def 	snake_case__			(		self)	->						List[Any]:

     """simple docstring"""



     _UpperCAmelCase			: Tuple     					=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_state_equivalency(*_A)


 def 	snake_case__			(		self)	->						Optional[int]:

     """simple docstring"""



     _UpperCAmelCase					,	_UpperCAmelCase			: str     					=			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         _UpperCAmelCase			: Dict     					=			model_class(config=_A)
         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[Any]:

     """simple docstring"""



     _UpperCAmelCase					,	_UpperCAmelCase			: Optional[Any]     					=			self.model_tester.prepare_config_and_inputs_for_common()
     _UpperCAmelCase			: List[str]     					=			True

     _UpperCAmelCase			: Union[str, Any]     					=			getattr(self.model_tester							,					'''seq_length'''							,					_A)

     for model_class in self.all_model_classes:
         _UpperCAmelCase			: Tuple     					=			True
         _UpperCAmelCase			: int     					=			False
         _UpperCAmelCase			: Any     					=			True
         _UpperCAmelCase			: Any     					=			model_class(_A)
         model.to(_A)
         model.eval()

         _UpperCAmelCase			: Dict     					=			self._prepare_for_class(_A							,					_A)
         _UpperCAmelCase			: List[str]     					=			inputs['''input_ids'''].shape[0]
         with torch.no_grad():
             _UpperCAmelCase			: List[str]     					=			model(**_A)
         _UpperCAmelCase			: Tuple     					=			outputs.attentions
         self.assertEqual(len(_A)							,					self.model_tester.num_hidden_layers)

         # check that output_attentions also work using config
         del inputs_dict["output_attentions"]
         _UpperCAmelCase			: int     					=			True
         _UpperCAmelCase			: Optional[int]     					=			model_class(_A)
         model.to(_A)
         model.eval()

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

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

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

         _UpperCAmelCase			: List[str]     					=			self._prepare_for_class(_A							,					_A)
         _UpperCAmelCase			: Dict     					=			inputs['''input_ids'''].shape[0]
         with torch.no_grad():
             _UpperCAmelCase			: Optional[int]     					=			model(**_A)

         _UpperCAmelCase			: Dict     					=			1
         self.assertEqual(out_len + added_hidden_states							,					len(_A))

         _UpperCAmelCase			: List[str]     					=			outputs.attentions

         self.assertEqual(len(_A)							,					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)	->						int:

     """simple docstring"""



     for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _UpperCAmelCase			: Optional[int]     					=			RwkvModel.from_pretrained(_A)
         self.assertIsNotNone(_A)



@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0	,				reason="See https://github.com/huggingface/transformers/pull/24204"     )
@slow
class 							A_      (  unittest.TestCase     ):






 '''simple docstring'''




 def 	snake_case__			(		self)	->						Optional[int]:

     """simple docstring"""



     _UpperCAmelCase			: List[Any]     					=			'''RWKV/rwkv-4-169m-pile'''
     _UpperCAmelCase			: Optional[Any]     					=			AutoTokenizer.from_pretrained(self.model_id)


 def 	snake_case__			(		self)	->						List[Any]:

     """simple docstring"""



     _UpperCAmelCase			: str     					=			'''Hello my name is Jasmine and I am a newbie to the'''
     _UpperCAmelCase			: Dict     					=			RwkvForCausalLM.from_pretrained(self.model_id).to(_A)

     _UpperCAmelCase			: Dict     					=			self.tokenizer('''Hello my name is'''							,					return_tensors='''pt''').input_ids.to(_A)
     _UpperCAmelCase			: int     					=			model.generate(_A							,					max_new_tokens=10)
     _UpperCAmelCase			: Optional[Any]     					=			self.tokenizer.decode(output[0].tolist())

     self.assertEqual(_A							,					_A)


 def 	snake_case__			(		self)	->						Optional[int]:

     """simple docstring"""



     _UpperCAmelCase			: Optional[int]     					=			'''Hello my name is Jasmine and I am a newbie to the'''

     _UpperCAmelCase			: Optional[int]     					=			self.tokenizer('''Hello my name is'''							,					return_tensors='''pt''').input_ids.to(_A)
     _UpperCAmelCase			: List[str]     					=			RwkvForCausalLM.from_pretrained(self.model_id							,					torch_dtype=torch.bfloataa).to(_A)

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

     self.assertEqual(_A							,					_A)












from unittest.mock import patch

import datasets
from datasets import Dataset

def        _lowerCamelCase		(  ) ->   Tuple:
    _UpperCAmelCase			: Union[str, Any]     					=			Dataset.from_dict({'''col_1''': [3, 2, 0, 1]}	)

    with patch('''tqdm.auto.tqdm'''	) as mock_tqdm:
        datasets.disable_progress_bar()
        dset.map(lambda __A	: {"col_2": x["col_1"] + 1}	)
        mock_tqdm.assert_not_called()

        mock_tqdm.reset_mock()

        datasets.enable_progress_bar()
        dset.map(lambda __A	: {"col_2": x["col_1"] + 1}	)
        mock_tqdm.assert_called()




