


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





     def __init__(self	,  UpperCAmelCase	,  UpperCAmelCase=1_4	,  UpperCAmelCase=7	,  UpperCAmelCase=True	,  UpperCAmelCase=False	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=9_9	,  UpperCAmelCase=3_2	,  UpperCAmelCase=5	,  UpperCAmelCase=3_7	,  UpperCAmelCase="gelu"	,  UpperCAmelCase=0.1	,  UpperCAmelCase=0.1	,  UpperCAmelCase=5_1_2	,  UpperCAmelCase=1_6	,  UpperCAmelCase=2	,  UpperCAmelCase=3	,  UpperCAmelCase=4	,  UpperCAmelCase=None	,  )  ->     Union[str, Any]:
      _lowercase							   =parent
      _lowercase							   =batch_size
      _lowercase							   =seq_length
      _lowercase							   =is_training
      _lowercase							   =use_token_type_ids
      _lowercase							   =use_input_mask
      _lowercase							   =use_labels
      _lowercase							   =use_mc_token_ids
      _lowercase							   =vocab_size
      _lowercase							   =hidden_size
      _lowercase							   =num_hidden_layers
      _lowercase							   =intermediate_size
      _lowercase							   =hidden_act
      _lowercase							   =hidden_dropout_prob
      _lowercase							   =attention_probs_dropout_prob
      _lowercase							   =max_position_embeddings
      _lowercase							   =type_vocab_size
      _lowercase							   =type_sequence_label_size
      _lowercase							   =num_labels
      _lowercase							   =num_choices
      _lowercase							   =scope
      _lowercase							   =vocab_size - 1
      _lowercase							   =vocab_size - 1
      _lowercase							   =vocab_size - 1





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





     def   __A      (self	,  UpperCAmelCase=False	,  UpperCAmelCase=False	,  UpperCAmelCase=False )  ->     List[str]:
      _lowercase							   =ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size )

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

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

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

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

      _lowercase							   =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   __A      (self	,  UpperCAmelCase=False	,  UpperCAmelCase=False	,  UpperCAmelCase=False )  ->     str:
      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   __A      (self )  ->     Optional[int]:
      _lowercase							   =self.get_config()
      _lowercase							   =3_0_0
      return config





     def   __A      (self )  ->     Optional[Any]:
      (
          (
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,										
      )							   =self.prepare_config_and_inputs()

      _lowercase							   =floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
      _lowercase							   =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   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  *UpperCAmelCase )  ->     str:
      _lowercase							   =True
      _lowercase							   =RwkvModel(config=UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()

      _lowercase							   =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   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  *UpperCAmelCase )  ->     str:
      _lowercase							   =RwkvForCausalLM(UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()

      _lowercase							   =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   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  *UpperCAmelCase )  ->     int:
      _lowercase							   =RwkvModel(config=UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()

      _lowercase							   =model(UpperCAmelCase )
      _lowercase							   =outputs.last_hidden_state

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

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

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





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  *UpperCAmelCase	,  UpperCAmelCase=False )  ->     List[Any]:
      _lowercase							   =RwkvForCausalLM(UpperCAmelCase )
      model.to(UpperCAmelCase )
      if gradient_checkpointing:
       model.gradient_checkpointing_enable()

      _lowercase							   =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   __A      (self )  ->     int:
      _lowercase							   =self.prepare_config_and_inputs()

      (
          (
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,					(
          _lowercase					
      )   ,										
      )							   =config_and_inputs

      _lowercase							   ={'''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):
     SCREAMING_SNAKE_CASE__       =    (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE__       =    (
         {'''feature-extraction''': RwkvModel, '''text-generation''': RwkvForCausalLM} if is_torch_available() else {}
     )
     # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE__       =    False
     SCREAMING_SNAKE_CASE__       =    False
     SCREAMING_SNAKE_CASE__       =    False
     SCREAMING_SNAKE_CASE__       =    False
     SCREAMING_SNAKE_CASE__       =    False  # Rwkv does not support head masking





     def   __A      (self )  ->     Optional[Any]:
      _lowercase							   =RwkvModelTester(self )
      _lowercase							   =ConfigTester(
          self	,  config_class=UpperCAmelCase	,  n_embd=3_7	,  common_properties=['''hidden_size''', '''num_hidden_layers'''] )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase=None )  ->     str:
      if isinstance(UpperCAmelCase	,  torch.Tensor ):
       _lowercase   ,					_lowercase							   =member.max().item(), member.min().item()
      elif isinstance(UpperCAmelCase	,  UpperCAmelCase ) or isinstance(UpperCAmelCase	,  UpperCAmelCase ):
       _lowercase   ,					_lowercase							   =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							   =container

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

      if not is_inside_interval:
       _lowercase							   ='''%s not found in %s''' % (safe_repr(UpperCAmelCase ), safe_repr(UpperCAmelCase ))
       self.fail(self._formatMessage(UpperCAmelCase	,  UpperCAmelCase ) )





     def   __A      (self )  ->     Optional[int]:
      self.config_tester.run_common_tests()





     def   __A      (self )  ->     Dict:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_rwkv_model(*UpperCAmelCase )





     def   __A      (self )  ->     Dict:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_causl_lm(*UpperCAmelCase )





     def   __A      (self )  ->     Union[str, Any]:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_state_equivalency(*UpperCAmelCase )





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

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

      _lowercase							   =getattr(self.model_tester	,  '''seq_length'''	,  UpperCAmelCase )

      for model_class in self.all_model_classes:
       _lowercase							   =True
       _lowercase							   =False
       _lowercase							   =True
       _lowercase							   =model_class(UpperCAmelCase )
       model.to(UpperCAmelCase )
       model.eval()

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

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

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

       _lowercase							   =self._prepare_for_class(UpperCAmelCase	,  UpperCAmelCase )
       _lowercase							   =inputs['''input_ids'''].shape[0]
       with torch.no_grad():
        _lowercase							   =model(**UpperCAmelCase )

       _lowercase							   =1
       self.assertEqual(out_len + added_hidden_states	,  len(UpperCAmelCase ) )

       _lowercase							   =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   __A      (self )  ->     int:
      for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _lowercase							   =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    lowerCamelCase__     ( unittest.TestCase):





     def   __A      (self )  ->     Optional[Any]:
      _lowercase							   ='''RWKV/rwkv-4-169m-pile'''
      _lowercase							   =AutoTokenizer.from_pretrained(self.model_id )





     def   __A      (self )  ->     str:
      _lowercase							   ='''Hello my name is Jasmine and I am a newbie to the'''
      _lowercase							   =RwkvForCausalLM.from_pretrained(self.model_id ).to(UpperCAmelCase )

      _lowercase							   =self.tokenizer('''Hello my name is'''	,  return_tensors='''pt''' ).input_ids.to(UpperCAmelCase )
      _lowercase							   =model.generate(UpperCAmelCase	,  max_new_tokens=1_0 )
      _lowercase							   =self.tokenizer.decode(output[0].tolist() )

      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Tuple:
      _lowercase							   ='''Hello my name is Jasmine and I am a newbie to the'''

      _lowercase							   =self.tokenizer('''Hello my name is'''	,  return_tensors='''pt''' ).input_ids.to(UpperCAmelCase )
      _lowercase							   =RwkvForCausalLM.from_pretrained(self.model_id	,  torch_dtype=torch.bfloataa ).to(UpperCAmelCase )

      _lowercase							   =model.generate(UpperCAmelCase	,  max_new_tokens=1_0 )
      _lowercase							   =self.tokenizer.decode(output[0].tolist() )

      self.assertEqual(UpperCAmelCase	,  UpperCAmelCase )





from ..utils import DummyObject, requires_backends







class    lowerCamelCase__     ( metaclass=lowerCAmelCase):
     SCREAMING_SNAKE_CASE__       =    ['''tensorflow_text''']





     def __init__(self	,  *UpperCAmelCase	,  **UpperCAmelCase )  ->     Any:
      requires_backends(self	,  ['''tensorflow_text'''] )

