




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




class       UpperCAmelCase_    :




    def __init__(       self		:			Any      ,					snake_case_		:			Optional[Any]      ,					snake_case_		:			Tuple=14      ,					snake_case_		:			Any=7      ,					snake_case_		:			Optional[int]=True      ,					snake_case_		:			Tuple=False      ,					snake_case_		:			Dict=True      ,					snake_case_		:			Any=True      ,					snake_case_		:			Dict=True      ,					snake_case_		:			Any=99      ,					snake_case_		:			Optional[Any]=32      ,					snake_case_		:			str=5      ,					snake_case_		:			List[str]=37      ,					snake_case_		:			Union[str, Any]="gelu"      ,					snake_case_		:			Optional[Any]=0.1      ,					snake_case_		:			List[Any]=0.1      ,					snake_case_		:			Optional[int]=512      ,					snake_case_		:			Tuple=16      ,					snake_case_		:			Any=2      ,					snake_case_		:			List[str]=3      ,					snake_case_		:			List[Any]=4      ,					snake_case_		:			Tuple=None      ,					) ->				Optional[Any]:




     '''simple docstring'''

     A__			=							parent
     A__			=							batch_size
     A__			=							seq_length
     A__			=							is_training
     A__			=							use_token_type_ids
     A__			=							use_input_mask
     A__			=							use_labels
     A__			=							use_mc_token_ids
     A__			=							vocab_size
     A__			=							hidden_size
     A__			=							num_hidden_layers
     A__			=							intermediate_size
     A__			=							hidden_act
     A__			=							hidden_dropout_prob
     A__			=							attention_probs_dropout_prob
     A__			=							max_position_embeddings
     A__			=							type_vocab_size
     A__			=							type_sequence_label_size
     A__			=							num_labels
     A__			=							num_choices
     A__			=							scope
     A__			=							vocab_size - 1
     A__			=							vocab_size - 1
     A__			=							vocab_size - 1




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




     '''simple docstring'''

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




    def 				__magic_name__				(       self		:			str      ,					snake_case_		:			int=False      ,					snake_case_		:			Optional[int]=False      ,					snake_case_		:			str=False ) ->				Union[str, Any]:




     '''simple docstring'''

     A__			=							ids_tensor([self.batch_size, self.seq_length]      ,					self.vocab_size )

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

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

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

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

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

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




    def 				__magic_name__				(       self		:			Optional[int]      ,					snake_case_		:			str=False      ,					snake_case_		:			List[Any]=False      ,					snake_case_		:			List[Any]=False ) ->				List[str]:




     '''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=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=snake_case_      ,					scale_attn_by_inverse_layer_idx=snake_case_      ,					reorder_and_upcast_attn=snake_case_      ,					)




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




     '''simple docstring'''

     A__			=							self.get_config()
     A__			=							300
     return config




    def 				__magic_name__				(       self		:			str ) ->				List[Any]:




     '''simple docstring'''

     (
         (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), 
     )			=							self.prepare_config_and_inputs()

     A__			=							floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
     A__			=							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 				__magic_name__				(       self		:			Optional[int]      ,					snake_case_		:			str      ,					snake_case_		:			Dict      ,					snake_case_		:			Union[str, Any]      ,					snake_case_		:			List[Any]      ,					snake_case_		:			str      ,					*snake_case_		:			Optional[int] ) ->				str:




     '''simple docstring'''

     A__			=							True
     A__			=							RwkvModel(config=snake_case_ )
     model.to(snake_case_ )
     model.eval()

     A__			=							model(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 				__magic_name__				(       self		:			Dict      ,					snake_case_		:			int      ,					snake_case_		:			List[str]      ,					snake_case_		:			List[Any]      ,					snake_case_		:			List[str]      ,					snake_case_		:			Union[str, Any]      ,					*snake_case_		:			List[Any] ) ->				Optional[Any]:




     '''simple docstring'''

     A__			=							RwkvForCausalLM(snake_case_ )
     model.to(snake_case_ )
     model.eval()

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




    def 				__magic_name__				(       self		:			str      ,					snake_case_		:			List[str]      ,					snake_case_		:			List[Any]      ,					snake_case_		:			List[str]      ,					snake_case_		:			List[Any]      ,					snake_case_		:			List[str]      ,					*snake_case_		:			int ) ->				Dict:




     '''simple docstring'''

     A__			=							RwkvModel(config=snake_case_ )
     model.to(snake_case_ )
     model.eval()

     A__			=							model(snake_case_ )
     A__			=							outputs.last_hidden_state

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

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

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




    def 				__magic_name__				(       self		:			Tuple      ,					snake_case_		:			str      ,					snake_case_		:			Tuple      ,					snake_case_		:			Any      ,					snake_case_		:			Optional[int]      ,					snake_case_		:			List[Any]      ,					*snake_case_		:			Union[str, Any]      ,					snake_case_		:			str=False ) ->				str:




     '''simple docstring'''

     A__			=							RwkvForCausalLM(snake_case_ )
     model.to(snake_case_ )
     if gradient_checkpointing:
      model.gradient_checkpointing_enable()

     A__			=							model(snake_case_      ,					labels=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 				__magic_name__				(       self		:			Tuple ) ->				Dict:




     '''simple docstring'''

     A__			=							self.prepare_config_and_inputs()

     (
         (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), (
         A__
     ), 
     )			=							config_and_inputs

     A__			=							{"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_    (		A_,  A_,  A_,  unittest.TestCase				):
    lowercase__		   = (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
    lowercase__		   = (
        {'''feature-extraction''': RwkvModel, '''text-generation''': RwkvForCausalLM} if is_torch_available() else {}
    )
    # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
    lowercase__		   = False
    lowercase__		   = False
    lowercase__		   = False
    lowercase__		   = False
    lowercase__		   = False  # Rwkv does not support head masking




    def 				__magic_name__				(       self		:			Any ) ->				str:




     '''simple docstring'''

     A__			=							RwkvModelTester(self )
     A__			=							ConfigTester(
         self      ,					config_class=snake_case_      ,					n_embd=37      ,					common_properties=["hidden_size", "num_hidden_layers"] )




    def 				__magic_name__				(       self		:			Dict      ,					snake_case_		:			Optional[Any]      ,					snake_case_		:			Dict      ,					snake_case_		:			Tuple=None ) ->				Optional[Any]:




     '''simple docstring'''

     if isinstance(snake_case_      ,					torch.Tensor ):
      A__, A__			=							member.max().item(), member.min().item()
     elif isinstance(snake_case_      ,					snake_case_ ) or isinstance(snake_case_      ,					snake_case_ ):
      A__, A__			=							max(snake_case_ ), min(snake_case_ )

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

     A__, A__			=							container

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

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




    def 				__magic_name__				(       self		:			int ) ->				int:




     '''simple docstring'''

     self.config_tester.run_common_tests()




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




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_rwkv_model(*snake_case_ )




    def 				__magic_name__				(       self		:			List[str] ) ->				Tuple:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_causl_lm(*snake_case_ )




    def 				__magic_name__				(       self		:			Tuple ) ->				List[str]:




     '''simple docstring'''

     A__			=							self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_state_equivalency(*snake_case_ )




    def 				__magic_name__				(       self		:			List[Any] ) ->				Tuple:




     '''simple docstring'''

     A__, A__			=							self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      A__			=							model_class(config=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 				__magic_name__				(       self		:			List[str] ) ->				Optional[Any]:




     '''simple docstring'''

     A__, A__			=							self.model_tester.prepare_config_and_inputs_for_common()
     A__			=							True

     A__			=							getattr(self.model_tester      ,					"seq_length"      ,					snake_case_ )

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

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

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

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

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

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

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

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

      A__			=							outputs.attentions

      self.assertEqual(len(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 				__magic_name__				(       self		:			Tuple ) ->				Optional[int]:




     '''simple docstring'''

     for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      A__			=							RwkvModel.from_pretrained(snake_case_ )
      self.assertIsNotNone(snake_case_ )






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




    def 				__magic_name__				(       self		:			Optional[Any] ) ->				str:




     '''simple docstring'''

     A__			=							"RWKV/rwkv-4-169m-pile"
     A__			=							AutoTokenizer.from_pretrained(self.model_id )




    def 				__magic_name__				(       self		:			int ) ->				Optional[Any]:




     '''simple docstring'''

     A__			=							"Hello my name is Jasmine and I am a newbie to the"
     A__			=							RwkvForCausalLM.from_pretrained(self.model_id ).to(snake_case_ )

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

     self.assertEqual(snake_case_      ,					snake_case_ )




    def 				__magic_name__				(       self		:			str ) ->				str:




     '''simple docstring'''

     A__			=							"Hello my name is Jasmine and I am a newbie to the"

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

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

     self.assertEqual(snake_case_      ,					snake_case_ )











"""simple docstring"""



def     _SCREAMING_SNAKE_CASE  (				lowercase_ = 10_00		)		->						int:
 A__			=							0
 A__			=							0
 while 1:
  num += 3
  if num >= n:
   break
  total += num
  num += 2
  if num >= n:
   break
  total += num
  num += 1
  if num >= n:
   break
  total += num
  num += 3
  if num >= n:
   break
  total += num
  num += 1
  if num >= n:
   break
  total += num
  num += 2
  if num >= n:
   break
  total += num
  num += 3
  if num >= n:
   break
  total += num
 return total


if __name__ == "__main__":
      print(f'{solution() = }')





