


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

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

       def     UpperCAmelCase( self   :  List[Any]     ):
            return RwkvConfig.from_pretrained("""sgugger/rwkv-4-pile-7b"""     )

       def     UpperCAmelCase( self   :  Union[str, Any]   ,				lowerCamelCase_   :  List[Any]=False   ,				lowerCamelCase_   :  Optional[int]=False   ,				lowerCamelCase_   :  Union[str, Any]=False     ):
            a_     :  Tuple				      =     ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size     )

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

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

            a_     :  Union[str, Any]				      =     None
            if self.use_mc_token_ids:
                 a_     :  List[Any]				      =     ids_tensor([self.batch_size, self.num_choices]   ,				self.seq_length     )

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

            a_     :  Optional[Any]				      =     self.get_config(
                gradient_checkpointing=lowerCamelCase_   ,				scale_attn_by_inverse_layer_idx=lowerCamelCase_   ,				reorder_and_upcast_attn=lowerCamelCase_   ,				)

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

       def     UpperCAmelCase( self   :  Any   ,				lowerCamelCase_   :  Union[str, Any]=False   ,				lowerCamelCase_   :  List[str]=False   ,				lowerCamelCase_   :  Union[str, Any]=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=lowerCamelCase_   ,				bos_token_id=self.bos_token_id   ,				eos_token_id=self.eos_token_id   ,				pad_token_id=self.pad_token_id   ,				gradient_checkpointing=lowerCamelCase_   ,				scale_attn_by_inverse_layer_idx=lowerCamelCase_   ,				reorder_and_upcast_attn=lowerCamelCase_   ,				)

       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  Tuple				      =     self.get_config()
            a_     :  int				      =     3_0_0
            return config

       def     UpperCAmelCase( self   :  List[Any]     ):
            (
                (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      
            )     :  Tuple				      =     self.prepare_config_and_inputs()

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

            a_     :  List[str]				      =     model(lowerCamelCase_     )

            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   :  Union[str, Any]   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  Tuple   ,				*lowerCamelCase_   :  List[Any]     ):
            a_     :  Optional[Any]				      =     RwkvForCausalLM(lowerCamelCase_     )
            model.to(lowerCamelCase_     )
            model.eval()

            a_     :  Tuple				      =     model(lowerCamelCase_   ,				labels=lowerCamelCase_     )
            self.parent.assertEqual(result.loss.shape   ,				()     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)     )

       def     UpperCAmelCase( self   :  Union[str, Any]   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  List[Any]   ,				*lowerCamelCase_   :  Union[str, Any]     ):
            a_     :  Union[str, Any]				      =     RwkvModel(config=lowerCamelCase_     )
            model.to(lowerCamelCase_     )
            model.eval()

            a_     :  Optional[Any]				      =     model(lowerCamelCase_     )
            a_     :  str				      =     outputs.last_hidden_state

            a_     :  int				      =     model(input_ids[:, :2]     )
            a_     :  Tuple				      =     outputs.last_hidden_state

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

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

       def     UpperCAmelCase( self   :  List[Any]   ,				lowerCamelCase_   :  str   ,				lowerCamelCase_   :  Tuple   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  List[Any]   ,				lowerCamelCase_   :  List[Any]   ,				*lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  int=False     ):
            a_     :  List[Any]				      =     RwkvForCausalLM(lowerCamelCase_     )
            model.to(lowerCamelCase_     )
            if gradient_checkpointing:
                 model.gradient_checkpointing_enable()

            a_     :  Any				      =     model(lowerCamelCase_   ,				labels=lowerCamelCase_     )
            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   :  List[str]     ):
            a_     :  Optional[int]				      =     self.prepare_config_and_inputs()

            (
                (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      (
                a_
            )      ,      
            )     :  Union[str, Any]				      =     config_and_inputs

            a_     :  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 	a__						(    lowerCAmelCase_    ,							lowerCAmelCase_    ,							lowerCAmelCase_    ,							unittest.TestCase     ):
       lowerCamelCase__:  Optional[Any]				     =  (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
       lowerCamelCase__:  List[str]				     =  (
           {"""feature-extraction""": RwkvModel, """text-generation""": RwkvForCausalLM} if is_torch_available() else {}
       )
       # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
       lowerCamelCase__:  Optional[int]				     =  False
       lowerCamelCase__:  Any				     =  False
       lowerCamelCase__:  Tuple				     =  False
       lowerCamelCase__:  Tuple				     =  False
       lowerCamelCase__:  int				     =  False  # Rwkv does not support head masking

       def     UpperCAmelCase( self   :  Any     ):
            a_     :  str				      =     RwkvModelTester(self     )
            a_     :  Dict				      =     ConfigTester(
                self   ,				config_class=lowerCamelCase_   ,				n_embd=3_7   ,				common_properties=["""hidden_size""", """num_hidden_layers"""]     )

       def     UpperCAmelCase( self   :  int   ,				lowerCamelCase_   :  Tuple   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  Dict=None     ):
            if isinstance(lowerCamelCase_   ,				torch.Tensor     ):
                 a_      ,      a_     :  Any				      =     member.max().item(), member.min().item()
            elif isinstance(lowerCamelCase_   ,				lowerCamelCase_     ) or isinstance(lowerCamelCase_   ,				lowerCamelCase_     ):
                 a_      ,      a_     :  List[str]				      =     max(lowerCamelCase_     ), min(lowerCamelCase_     )

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

            a_      ,      a_     :  Any				      =     container

            a_     :  Dict				      =     (min_value >= expected_min) and (max_value <= expected_max)

            if not is_inside_interval:
                 a_     :  Any				      =     """%s not found in %s""" % (safe_repr(lowerCamelCase_     ), safe_repr(lowerCamelCase_     ))
                 self.fail(self._formatMessage(lowerCamelCase_   ,				lowerCamelCase_     )     )

       def     UpperCAmelCase( self   :  Any     ):
            self.config_tester.run_common_tests()

       def     UpperCAmelCase( self   :  Optional[int]     ):
            a_     :  List[Any]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_rwkv_model(*lowerCamelCase_     )

       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  Optional[int]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_causl_lm(*lowerCamelCase_     )

       def     UpperCAmelCase( self   :  Optional[Any]     ):
            a_     :  List[str]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_state_equivalency(*lowerCamelCase_     )

       def     UpperCAmelCase( self   :  int     ):
            a_      ,      a_     :  Optional[Any]				      =     self.model_tester.prepare_config_and_inputs_for_common()

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

            a_     :  int				      =     getattr(self.model_tester   ,				"""seq_length"""   ,				lowerCamelCase_     )

            for model_class in self.all_model_classes:
                 a_     :  Optional[int]				      =     True
                 a_     :  str				      =     False
                 a_     :  str				      =     True
                 a_     :  Union[str, Any]				      =     model_class(lowerCamelCase_     )
                 model.to(lowerCamelCase_     )
                 model.eval()

                 a_     :  int				      =     self._prepare_for_class(lowerCamelCase_   ,				lowerCamelCase_     )
                 a_     :  List[str]				      =     inputs["""input_ids"""].shape[0]
                 with torch.no_grad():
                      a_     :  int				      =     model(**lowerCamelCase_     )
                 a_     :  int				      =     outputs.attentions
                 self.assertEqual(len(lowerCamelCase_     )   ,				self.model_tester.num_hidden_layers     )

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

                 a_     :  List[str]				      =     self._prepare_for_class(lowerCamelCase_   ,				lowerCamelCase_     )
                 a_     :  Optional[int]				      =     inputs["""input_ids"""].shape[0]
                 with torch.no_grad():
                      a_     :  Dict				      =     model(**lowerCamelCase_     )
                 a_     :  List[str]				      =     outputs.attentions
                 self.assertEqual(len(lowerCamelCase_     )   ,				self.model_tester.num_hidden_layers     )

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

                 # Check attention is always last and order is fine
                 a_     :  str				      =     True
                 a_     :  Tuple				      =     True
                 a_     :  Dict				      =     model_class(lowerCamelCase_     )
                 model.to(lowerCamelCase_     )
                 model.eval()

                 a_     :  List[Any]				      =     self._prepare_for_class(lowerCamelCase_   ,				lowerCamelCase_     )
                 a_     :  int				      =     inputs["""input_ids"""].shape[0]
                 with torch.no_grad():
                      a_     :  Tuple				      =     model(**lowerCamelCase_     )

                 a_     :  int				      =     1
                 self.assertEqual(out_len + added_hidden_states   ,				len(lowerCamelCase_     )     )

                 a_     :  List[str]				      =     outputs.attentions

                 self.assertEqual(len(lowerCamelCase_     )   ,				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   :  Union[str, Any]     ):
            for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 a_     :  Dict				      =     RwkvModel.from_pretrained(lowerCamelCase_     )
                 self.assertIsNotNone(lowerCamelCase_     )


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

       def     UpperCAmelCase( self   :  str     ):
            a_     :  Optional[int]				      =     """RWKV/rwkv-4-169m-pile"""
            a_     :  int				      =     AutoTokenizer.from_pretrained(self.model_id     )

       def     UpperCAmelCase( self   :  Optional[int]     ):
            a_     :  Optional[int]				      =     """Hello my name is Jasmine and I am a newbie to the"""
            a_     :  List[Any]				      =     RwkvForCausalLM.from_pretrained(self.model_id     ).to(lowerCamelCase_     )

            a_     :  str				      =     self.tokenizer("""Hello my name is"""   ,				return_tensors="""pt"""     ).input_ids.to(lowerCamelCase_     )
            a_     :  Optional[int]				      =     model.generate(lowerCamelCase_   ,				max_new_tokens=1_0     )
            a_     :  Tuple				      =     self.tokenizer.decode(output[0].tolist()     )

            self.assertEqual(lowerCamelCase_   ,				lowerCamelCase_     )

       def     UpperCAmelCase( self   :  Optional[Any]     ):
            a_     :  List[str]				      =     """Hello my name is Jasmine and I am a newbie to the"""

            a_     :  Dict				      =     self.tokenizer("""Hello my name is"""   ,				return_tensors="""pt"""     ).input_ids.to(lowerCamelCase_     )
            a_     :  Tuple				      =     RwkvForCausalLM.from_pretrained(self.model_id   ,				torch_dtype=torch.bfloataa     ).to(lowerCamelCase_     )

            a_     :  Optional[int]				      =     model.generate(lowerCamelCase_   ,				max_new_tokens=1_0     )
            a_     :  int				      =     self.tokenizer.decode(output[0].tolist()     )

            self.assertEqual(lowerCamelCase_   ,				lowerCamelCase_     )





'''simple docstring'''







import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
						import torch

						from transformers import GPTaLMHeadModel


@require_torch
@require_sentencepiece
@require_tokenizers
class       lowerCAmelCase			(					unittest.TestCase      ):




		@slow
		@unittest.skip("Model is not available."   )
		def      lowercase							(					self   ):
								_SCREAMING_SNAKE_CASE					=			"nvidia/megatron-gpt2-345m/"
								if "MYDIR" in os.environ:
														_SCREAMING_SNAKE_CASE					=			os.path.join(os.environ["MYDIR"]       ,    UpperCamelCase   )
								_SCREAMING_SNAKE_CASE					=			GPTaLMHeadModel.from_pretrained(UpperCamelCase   )
								model.to(UpperCamelCase   )
								model.half()

								_SCREAMING_SNAKE_CASE					=			torch.tensor(
								    [[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]       ,    device=UpperCamelCase       ,    dtype=torch.long       ,    )

								with torch.no_grad():
														_SCREAMING_SNAKE_CASE					=			model(UpperCamelCase   ).logits

								_SCREAMING_SNAKE_CASE					=			torch.Size((1, 9, 50_257)   )
								self.assertEqual(output.shape       ,    UpperCamelCase   )

								_SCREAMING_SNAKE_CASE					=			torch.tensor(
								    [
								        4.94_14,
								        -0.29_20,
								        -1.21_48,
								        -4.02_73,
								        -0.51_61,
								        -5.21_09,
								        -1.24_12,
								        -1.83_01,
								        -1.77_34,
								        -4.71_48,
								        -0.23_17,
								        -1.08_11,
								        -2.17_77,
								        0.41_41,
								        -3.79_69,
								        -4.05_86,
								        -2.53_32,
								        -3.38_09,
								        4.38_67,
								    ]       ,    device=UpperCamelCase       ,    dtype=torch.half       ,    )

								for i in range(19   ):
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			8 * i // 17, 2_792 * i  # along the diagonal
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			output[0, r, c], expected_diag[i]
														_SCREAMING_SNAKE_CASE					=			F'row={r} col={c} computed={computed} expected={expected}'
														self.assertAlmostEqual(UpperCamelCase       ,    UpperCamelCase       ,    delta=1e-4       ,    msg=UpperCamelCase   )