


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:
   UpperCAmelCase     =       False







class lowerCAmelCase_		:

 '''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=5_12						, _UpperCAmelCase=16						, _UpperCAmelCase=2						, _UpperCAmelCase=3						, _UpperCAmelCase=4						, _UpperCAmelCase=None						, ):
       snake_case_			= parent
       snake_case_			= batch_size
       snake_case_			= seq_length
       snake_case_			= is_training
       snake_case_			= use_token_type_ids
       snake_case_			= use_input_mask
       snake_case_			= use_labels
       snake_case_			= use_mc_token_ids
       snake_case_			= vocab_size
       snake_case_			= hidden_size
       snake_case_			= num_hidden_layers
       snake_case_			= intermediate_size
       snake_case_			= hidden_act
       snake_case_			= hidden_dropout_prob
       snake_case_			= attention_probs_dropout_prob
       snake_case_			= max_position_embeddings
       snake_case_			= type_vocab_size
       snake_case_			= type_sequence_label_size
       snake_case_			= num_labels
       snake_case_			= num_choices
       snake_case_			= scope
       snake_case_			= vocab_size - 1
       snake_case_			= vocab_size - 1
       snake_case_			= vocab_size - 1



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



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase=False						, _UpperCAmelCase=False						, _UpperCAmelCase=False   ):
       snake_case_			= ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size   )

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

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

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

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

       snake_case_			= self.get_config(
           gradient_checkpointing=_SCREAMING_SNAKE_CASE						, scale_attn_by_inverse_layer_idx=_SCREAMING_SNAKE_CASE						, reorder_and_upcast_attn=_SCREAMING_SNAKE_CASE						, )

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



 def 		UpperCamelCase__	(      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=_SCREAMING_SNAKE_CASE						, bos_token_id=self.bos_token_id						, eos_token_id=self.eos_token_id						, pad_token_id=self.pad_token_id						, gradient_checkpointing=_SCREAMING_SNAKE_CASE						, scale_attn_by_inverse_layer_idx=_SCREAMING_SNAKE_CASE						, reorder_and_upcast_attn=_SCREAMING_SNAKE_CASE						, )



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= self.get_config()
       snake_case_			= 3_00
       return config



 def 		UpperCamelCase__	(      self   ):
       (
           (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       
       )			= self.prepare_config_and_inputs()

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

       snake_case_			= model(_SCREAMING_SNAKE_CASE   )

       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 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, *_UpperCAmelCase   ):
       snake_case_			= RwkvForCausalLM(_SCREAMING_SNAKE_CASE   )
       model.to(_SCREAMING_SNAKE_CASE   )
       model.eval()

       snake_case_			= model(_SCREAMING_SNAKE_CASE						, labels=_SCREAMING_SNAKE_CASE   )
       self.parent.assertEqual(result.loss.shape						, ()   )
       self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, *_UpperCAmelCase   ):
       snake_case_			= RwkvModel(config=_SCREAMING_SNAKE_CASE   )
       model.to(_SCREAMING_SNAKE_CASE   )
       model.eval()

       snake_case_			= model(_SCREAMING_SNAKE_CASE   )
       snake_case_			= outputs.last_hidden_state

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

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

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



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, *_UpperCAmelCase						, _UpperCAmelCase=False   ):
       snake_case_			= RwkvForCausalLM(_SCREAMING_SNAKE_CASE   )
       model.to(_SCREAMING_SNAKE_CASE   )
       if gradient_checkpointing:
             model.gradient_checkpointing_enable()

       snake_case_			= model(_SCREAMING_SNAKE_CASE						, labels=_SCREAMING_SNAKE_CASE   )
       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 		UpperCamelCase__	(      self   ):
       snake_case_			= self.prepare_config_and_inputs()

       (
           (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       (
           snake_case_
       )			,       
       )			= config_and_inputs

       snake_case_			= {'''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_		(					_UpperCamelCase   ,		_UpperCamelCase   ,		_UpperCamelCase   ,		unittest.TestCase			):

 '''simple docstring'''

 __snake_case				=    (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
 __snake_case				=    (
     {'feature-extraction': RwkvModel, 'text-generation': RwkvForCausalLM} if is_torch_available() else {}
 )
 # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
 __snake_case				=    False
 __snake_case				=    False
 __snake_case				=    False
 __snake_case				=    False
 __snake_case				=    False  # Rwkv does not support head masking



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= RwkvModelTester(self   )
       snake_case_			= ConfigTester(
           self						, config_class=_SCREAMING_SNAKE_CASE						, n_embd=37						, common_properties=['''hidden_size''', '''num_hidden_layers''']   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase=None   ):
       if isinstance(_SCREAMING_SNAKE_CASE						, torch.Tensor   ):
             snake_case_			,       snake_case_			= member.max().item(), member.min().item()
       elif isinstance(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE   ) or isinstance(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE   ):
             snake_case_			,       snake_case_			= max(_SCREAMING_SNAKE_CASE   ), min(_SCREAMING_SNAKE_CASE   )

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

       snake_case_			,       snake_case_			= container

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

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



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



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_rwkv_model(*_SCREAMING_SNAKE_CASE   )



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_causl_lm(*_SCREAMING_SNAKE_CASE   )



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_state_equivalency(*_SCREAMING_SNAKE_CASE   )



 def 		UpperCamelCase__	(      self   ):
       snake_case_			,       snake_case_			= self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             snake_case_			= model_class(config=_SCREAMING_SNAKE_CASE   )
             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 		UpperCamelCase__	(      self   ):
       snake_case_			,       snake_case_			= self.model_tester.prepare_config_and_inputs_for_common()
       snake_case_			= True

       snake_case_			= getattr(self.model_tester						, '''seq_length'''						, _SCREAMING_SNAKE_CASE   )

       for model_class in self.all_model_classes:
             snake_case_			= True
             snake_case_			= False
             snake_case_			= True
             snake_case_			= model_class(_SCREAMING_SNAKE_CASE   )
             model.to(_SCREAMING_SNAKE_CASE   )
             model.eval()

             snake_case_			= self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE   )
             snake_case_			= inputs['''input_ids'''].shape[0]
             with torch.no_grad():
                   snake_case_			= model(**_SCREAMING_SNAKE_CASE   )
             snake_case_			= outputs.attentions
             self.assertEqual(len(_SCREAMING_SNAKE_CASE   )						, self.model_tester.num_hidden_layers   )

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             snake_case_			= True
             snake_case_			= model_class(_SCREAMING_SNAKE_CASE   )
             model.to(_SCREAMING_SNAKE_CASE   )
             model.eval()

             snake_case_			= self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE   )
             snake_case_			= inputs['''input_ids'''].shape[0]
             with torch.no_grad():
                   snake_case_			= model(**_SCREAMING_SNAKE_CASE   )
             snake_case_			= outputs.attentions
             self.assertEqual(len(_SCREAMING_SNAKE_CASE   )						, self.model_tester.num_hidden_layers   )

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

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

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

             snake_case_			= 1
             self.assertEqual(out_len + added_hidden_states						, len(_SCREAMING_SNAKE_CASE   )   )

             snake_case_			= outputs.attentions

             self.assertEqual(len(_SCREAMING_SNAKE_CASE   )						, self.model_tester.num_hidden_layers   )
             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]   )						, [batch_size, seq_len, config.hidden_size]						, )





 @slow
 def 		UpperCamelCase__	(      self   ):
       for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_			= RwkvModel.from_pretrained(_SCREAMING_SNAKE_CASE   )
             self.assertIsNotNone(_SCREAMING_SNAKE_CASE   )







@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			):

 '''simple docstring'''

 def 		UpperCamelCase__	(      self   ):
       snake_case_			= '''RWKV/rwkv-4-169m-pile'''
       snake_case_			= AutoTokenizer.from_pretrained(self.model_id   )



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= '''Hello my name is Jasmine and I am a newbie to the'''
       snake_case_			= RwkvForCausalLM.from_pretrained(self.model_id   ).to(_SCREAMING_SNAKE_CASE   )

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

       self.assertEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE   )





 def 		UpperCamelCase__	(      self   ):
       snake_case_			= '''Hello my name is Jasmine and I am a newbie to the'''

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

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

       self.assertEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE   )
'''simple docstring'''




def 	_UpperCAmelCase       (				_UpperCamelCase       :      list[list[int]],  _UpperCamelCase       :      list[list[int]]  )      -> list[list[int]]:
    A_      =  []
    A_      =  len(_UpperCamelCase  )
    for i in range(_UpperCamelCase  ):
        A_      =  []
        for j in range(_UpperCamelCase  ):
            A_      =  0
            for k in range(_UpperCamelCase  ):
                A_      =  val + matrix_a[i][k] * matrix_b[k][j]
            list_a.append(_UpperCamelCase  )
        matrix_c.append(_UpperCamelCase  )
    return matrix_c





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int  )      -> list[list[int]]:
    return [[int(row == column  ) for column in range(_UpperCamelCase  )] for row in range(_UpperCamelCase  )]





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int  )      -> int:
    if n <= 1:
        return n
    A_      =  identity(2  )
    A_      =  [[1, 1], [1, 0]]
    A_      =  n - 1
    while n > 0:
        if n % 2 == 1:
            A_      =  multiply(_UpperCamelCase,  _UpperCamelCase  )
        A_      =  multiply(_UpperCamelCase,  _UpperCamelCase  )
        A_      =  int(n / 2  )
    return res_matrix[0][0]





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int  )      -> int:
    if n <= 1:
        return n
    A_      =  0
    A_      =  1
    for _ in range(2,  n + 1  ):
        A_  ,A_      =  fiba, fiba + fiba
    return fiba





def 	_UpperCAmelCase       (				)      -> None:
    for ordinal in "0th 1st 2nd 3rd 10th 100th 1000th".split():
        A_      =  int(''''''.join(c for c in ordinal if c in '''0123456789'''  )  )  # 1000th --> 1000
        print(
            F'''{ordinal} fibonacci number using matrix exponentiation is '''
            F'''{nth_fibonacci_matrix(_UpperCamelCase  )} and using bruteforce is '''
            F'''{nth_fibonacci_bruteforce(_UpperCamelCase  )}\n'''  )
    # from timeit import timeit
    # print(timeit("nth_fibonacci_matrix(1000000)",
    #              "from main import nth_fibonacci_matrix", number=5))
    # print(timeit("nth_fibonacci_bruteforce(1000000)",
    #              "from main import nth_fibonacci_bruteforce", number=5))
    # 2.3342058970001744
    # 57.256506615000035


if __name__ == "__main__":
  import doctest

  doctest.testmod()
  main()






