"""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 					UpperCAmelCase		:



    """simple docstring"""

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






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






    def _snake_case    (     self					,    _UpperCAmelCase=False					,    _UpperCAmelCase=False					,    _UpperCAmelCase=False							):
       lowercase__:     int      =   ids_tensor([self.batch_size, self.seq_length]					,    self.vocab_size							)

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

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

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

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

       lowercase__:     str      =   self.get_config(
           gradient_checkpointing=_UpperCAmelCase					,    scale_attn_by_inverse_layer_idx=_UpperCAmelCase					,    reorder_and_upcast_attn=_UpperCAmelCase					,    )

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






    def _snake_case    (     self					,    _UpperCAmelCase=False					,    _UpperCAmelCase=False					,    _UpperCAmelCase=False							):
       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=_UpperCAmelCase					,    bos_token_id=self.bos_token_id					,    eos_token_id=self.eos_token_id					,    pad_token_id=self.pad_token_id					,    gradient_checkpointing=_UpperCAmelCase					,    scale_attn_by_inverse_layer_idx=_UpperCAmelCase					,    reorder_and_upcast_attn=_UpperCAmelCase					,    )






    def _snake_case    (     self							):
       lowercase__:     List[str]      =   self.get_config()
       lowercase__:     Dict      =   300
       return config






    def _snake_case    (     self							):
       (
           (
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				
       ):     Any      =   self.prepare_config_and_inputs()

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

       lowercase__:     int      =   model(_UpperCAmelCase							)

       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					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    *_UpperCAmelCase							):
       lowercase__:     str      =   RwkvForCausalLM(_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()

       lowercase__:     Optional[Any]      =   model(_UpperCAmelCase					,    labels=_UpperCAmelCase							)
       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					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    *_UpperCAmelCase							):
       lowercase__:     Optional[Any]      =   RwkvModel(config=_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()

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

       lowercase__:     List[str]      =   model(input_ids[:, :2]							)
       lowercase__:     Tuple      =   outputs.last_hidden_state

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

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






    def _snake_case    (     self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    *_UpperCAmelCase					,    _UpperCAmelCase=False							):
       lowercase__:     Optional[Any]      =   RwkvForCausalLM(_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       if gradient_checkpointing:
          model.gradient_checkpointing_enable()

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

       (
           (
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				(
           lowercase__
       ),				
       ):     Optional[int]      =   config_and_inputs

       lowercase__:     List[str]      =   {'''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		(_UpperCAmelCase					,_UpperCAmelCase					,_UpperCAmelCase					,unittest.TestCase						):



    """simple docstring"""

    _UpperCAmelCase			:int							=				(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			:List[str]							=				False
    _UpperCAmelCase			:int							=				False
    _UpperCAmelCase			:List[str]							=				False
    _UpperCAmelCase			:List[str]							=				False
    _UpperCAmelCase			:Union[str, Any]							=				False  # Rwkv does not support head masking






    def _snake_case    (     self							):
       lowercase__:     str      =   RwkvModelTester(self							)
       lowercase__:     Union[str, Any]      =   ConfigTester(
           self					,    config_class=_UpperCAmelCase					,    n_embd=37					,    common_properties=['''hidden_size''', '''num_hidden_layers''']							)






    def _snake_case    (     self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase=None							):
       if isinstance(_UpperCAmelCase					,    torch.Tensor							):
          lowercase__,				lowercase__:     Optional[int]      =   member.max().item(), member.min().item()
       elif isinstance(_UpperCAmelCase					,    _UpperCAmelCase							) or isinstance(_UpperCAmelCase					,    _UpperCAmelCase							):
          lowercase__,				lowercase__:     str      =   max(_UpperCAmelCase							), min(_UpperCAmelCase							)

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

       lowercase__,				lowercase__:     Optional[Any]      =   container

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

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






    def _snake_case    (     self							):
       self.config_tester.run_common_tests()






    def _snake_case    (     self							):
       lowercase__:     List[str]      =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_rwkv_model(*_UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__:     Tuple      =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_causl_lm(*_UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__:     Dict      =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_state_equivalency(*_UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__,				lowercase__:     Optional[int]      =   self.model_tester.prepare_config_and_inputs_for_common()

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

       lowercase__:     Tuple      =   getattr(self.model_tester					,    '''seq_length'''					,    _UpperCAmelCase							)

       for model_class in self.all_model_classes:
          lowercase__:     int      =   True
          lowercase__:     List[Any]      =   False
          lowercase__:     Tuple      =   True
          lowercase__:     int      =   model_class(_UpperCAmelCase							)
          model.to(_UpperCAmelCase							)
          model.eval()

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

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

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

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

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

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

          lowercase__:     Tuple      =   1
          self.assertEqual(out_len + added_hidden_states					,    len(_UpperCAmelCase							)							)

          lowercase__:     Optional[int]      =   outputs.attentions

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



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






    def _snake_case    (     self							):
       lowercase__:     List[str]      =   '''Hello my name is Jasmine and I am a newbie to the'''
       lowercase__:     Optional[Any]      =   RwkvForCausalLM.from_pretrained(self.model_id							).to(_UpperCAmelCase							)

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

       self.assertEqual(_UpperCAmelCase					,    _UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__:     int      =   '''Hello my name is Jasmine and I am a newbie to the'''

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

       lowercase__:     int      =   model.generate(_UpperCAmelCase					,    max_new_tokens=10							)
       lowercase__:     Union[str, Any]      =   self.tokenizer.decode(output[0].tolist()							)

       self.assertEqual(_UpperCAmelCase					,    _UpperCAmelCase							)

"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging


__A         =				logging.get_logger(__name__)

__A         =				{
    "uw-madison/yoso-4096": "https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json",
    # See all YOSO models at https://huggingface.co/models?filter=yoso
}


class 					UpperCAmelCase		(_UpperCAmelCase						):



    """simple docstring"""

    _UpperCAmelCase			:Tuple							=				"yoso"






    def __init__(     self					,    _UpperCAmelCase=50265					,    _UpperCAmelCase=768					,    _UpperCAmelCase=12					,    _UpperCAmelCase=12					,    _UpperCAmelCase=3072					,    _UpperCAmelCase="gelu"					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=4096					,    _UpperCAmelCase=1					,    _UpperCAmelCase=0.02					,    _UpperCAmelCase=1e-1_2					,    _UpperCAmelCase="absolute"					,    _UpperCAmelCase=True					,    _UpperCAmelCase=9					,    _UpperCAmelCase=64					,    _UpperCAmelCase=None					,    _UpperCAmelCase=True					,    _UpperCAmelCase=True					,    _UpperCAmelCase=1					,    _UpperCAmelCase=0					,    _UpperCAmelCase=2					,    **_UpperCAmelCase					,    ):
       super().__init__(pad_token_id=_UpperCAmelCase					,    bos_token_id=_UpperCAmelCase					,    eos_token_id=_UpperCAmelCase					,    **_UpperCAmelCase							)

       lowercase__:     int      =   vocab_size
       lowercase__:     List[str]      =   max_position_embeddings
       lowercase__:     Tuple      =   hidden_size
       lowercase__:     Optional[Any]      =   num_hidden_layers
       lowercase__:     Optional[int]      =   num_attention_heads
       lowercase__:     List[Any]      =   intermediate_size
       lowercase__:     Union[str, Any]      =   hidden_act
       lowercase__:     int      =   hidden_dropout_prob
       lowercase__:     Optional[Any]      =   attention_probs_dropout_prob
       lowercase__:     List[Any]      =   initializer_range
       lowercase__:     Optional[Any]      =   type_vocab_size
       lowercase__:     int      =   layer_norm_eps
       lowercase__:     Optional[Any]      =   position_embedding_type
       lowercase__:     int      =   use_expectation
       lowercase__:     List[str]      =   hash_code_len
       lowercase__:     Union[str, Any]      =   num_hash
       lowercase__:     Any      =   conv_window
       lowercase__:     Dict      =   use_fast_hash
       lowercase__:     Dict      =   lsh_backward
