import copy
from typing import Dict

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING
from ..bit import BitConfig


UpperCAmelCase							   =      logging.get_logger(__name__)

UpperCAmelCase							   =      {
    """google/vit-hybrid-base-bit-384""": """https://huggingface.co/vit-hybrid-base-bit-384/resolve/main/config.json""",
    # See all ViT hybrid models at https://huggingface.co/models?filter=vit
}






class 							_SCREAMING_SNAKE_CASE(   snake_case__	):
   UpperCamelCase_				:						Union[str, Any] 			=   'vit-hybrid'
   def __init__(    self							,     lowerCAmelCase_=None							,     lowerCAmelCase_=7_68							,     lowerCAmelCase_=12							,     lowerCAmelCase_=12							,     lowerCAmelCase_=30_72							,     lowerCAmelCase_="gelu"							,     lowerCAmelCase_=0.0							,     lowerCAmelCase_=0.0							,     lowerCAmelCase_=0.02							,     lowerCAmelCase_=1e-12							,     lowerCAmelCase_=2_24							,     lowerCAmelCase_=1							,     lowerCAmelCase_=3							,     lowerCAmelCase_=[1, 10_24, 24, 24]							,     lowerCAmelCase_=True							,     **lowerCAmelCase_							,     )	->Union[str, Any]:
        super().__init__(**lowerCAmelCase_    )

        if backbone_config is None:
             logger.info('''`backbone_config` is `None`. Initializing the config with a `BiT` backbone.'''    )
             a  :    List[Any]				=						{
                 '''global_padding''': '''same''',
                 '''layer_type''': '''bottleneck''',
                 '''depths''': [3, 4, 9],
                 '''out_features''': ['''stage3'''],
                 '''embedding_dynamic_padding''': True,
             }

        if isinstance(lowerCAmelCase_							,     lowerCAmelCase_    ):
             if "model_type" in backbone_config:
                  a  :    List[str]				=						CONFIG_MAPPING[backbone_config['''model_type''']]
             else:
                  logger.info(
                      '''`model_type` is not found in `backbone_config`. Use `Bit` as the backbone configuration class.'''    )
                  a  :    Optional[Any]				=						BitConfig
             a  :    List[Any]				=						backbone_config_class(**lowerCAmelCase_    )

        a  :    Any				=						backbone_featmap_shape
        a  :    str				=						backbone_config
        a  :    Tuple				=						hidden_size
        a  :    List[str]				=						num_hidden_layers
        a  :    Tuple				=						num_attention_heads
        a  :    List[Any]				=						intermediate_size
        a  :    Optional[int]				=						hidden_act
        a  :    Optional[Any]				=						hidden_dropout_prob
        a  :    int				=						attention_probs_dropout_prob
        a  :    Optional[int]				=						initializer_range
        a  :    Optional[int]				=						layer_norm_eps
        a  :    Optional[Any]				=						image_size
        a  :    Union[str, Any]				=						patch_size
        a  :    List[Any]				=						num_channels
        a  :    Dict				=						qkv_bias


   def _lowercase(    self    )	->Dict[str, any]:
        a  :    Dict				=						copy.deepcopy(self.__dict__    )
        a  :    Any				=						self.backbone_config.to_dict()
        a  :    Tuple				=						self.__class__.model_type
        return output
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 							_SCREAMING_SNAKE_CASE:
   def __init__(    self							,     lowerCAmelCase_							,     lowerCAmelCase_=14							,     lowerCAmelCase_=7							,     lowerCAmelCase_=True							,     lowerCAmelCase_=False							,     lowerCAmelCase_=True							,     lowerCAmelCase_=True							,     lowerCAmelCase_=True							,     lowerCAmelCase_=99							,     lowerCAmelCase_=32							,     lowerCAmelCase_=5							,     lowerCAmelCase_=37							,     lowerCAmelCase_="gelu"							,     lowerCAmelCase_=0.1							,     lowerCAmelCase_=0.1							,     lowerCAmelCase_=5_12							,     lowerCAmelCase_=16							,     lowerCAmelCase_=2							,     lowerCAmelCase_=3							,     lowerCAmelCase_=4							,     lowerCAmelCase_=None							,     )	->int:
        a  :    Optional[int]				=						parent
        a  :    Optional[Any]				=						batch_size
        a  :    Union[str, Any]				=						seq_length
        a  :    int				=						is_training
        a  :    List[Any]				=						use_token_type_ids
        a  :    Tuple				=						use_input_mask
        a  :    Optional[int]				=						use_labels
        a  :    List[str]				=						use_mc_token_ids
        a  :    str				=						vocab_size
        a  :    Any				=						hidden_size
        a  :    List[Any]				=						num_hidden_layers
        a  :    Optional[int]				=						intermediate_size
        a  :    Dict				=						hidden_act
        a  :    Tuple				=						hidden_dropout_prob
        a  :    str				=						attention_probs_dropout_prob
        a  :    Optional[Any]				=						max_position_embeddings
        a  :    Optional[int]				=						type_vocab_size
        a  :    Any				=						type_sequence_label_size
        a  :    Any				=						num_labels
        a  :    int				=						num_choices
        a  :    List[str]				=						scope
        a  :    List[Any]				=						vocab_size - 1
        a  :    Optional[Any]				=						vocab_size - 1
        a  :    Optional[Any]				=						vocab_size - 1
   def _lowercase(    self    )	->Optional[int]:
        return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''    )
   def _lowercase(    self							,     lowerCAmelCase_=False							,     lowerCAmelCase_=False							,     lowerCAmelCase_=False    )	->Tuple:
        a  :    int				=						ids_tensor([self.batch_size, self.seq_length]							,     self.vocab_size    )

        a  :    List[str]				=						None
        if self.use_input_mask:
             a  :    Optional[int]				=						random_attention_mask([self.batch_size, self.seq_length]    )

        a  :    Union[str, Any]				=						None
        if self.use_token_type_ids:
             a  :    str				=						ids_tensor([self.batch_size, self.seq_length]							,     self.type_vocab_size    )

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

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

        a  :    Optional[int]				=						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 _lowercase(    self							,     lowerCAmelCase_=False							,     lowerCAmelCase_=False							,     lowerCAmelCase_=False    )	->Dict:
        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 _lowercase(    self    )	->Optional[Any]:
        a  :    Optional[Any]				=						self.get_config()
        a  :    Any				=						3_00
        return config
   def _lowercase(    self    )	->int:
        (
            (
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							
        )  :    Any				=						self.prepare_config_and_inputs()

        a  :    str				=						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 _lowercase(    self							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     *lowerCAmelCase_    )	->Optional[int]:
        a  :    List[str]				=						True
        a  :    Union[str, Any]				=						RwkvModel(config=lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()

        a  :    Tuple				=						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 _lowercase(    self							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     *lowerCAmelCase_    )	->Optional[int]:
        a  :    List[Any]				=						RwkvForCausalLM(lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()

        a  :    List[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)    )
   def _lowercase(    self							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     *lowerCAmelCase_    )	->Tuple:
        a  :    int				=						RwkvModel(config=lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()

        a  :    Dict				=						model(lowerCAmelCase_    )
        a  :    Union[str, Any]				=						outputs.last_hidden_state

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

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

        self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]							,     dim=1    )							,     lowerCAmelCase_							,     atol=1e-5    )    )
   def _lowercase(    self							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_							,     *lowerCAmelCase_							,     lowerCAmelCase_=False    )	->Dict:
        a  :    Dict				=						RwkvForCausalLM(lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        if gradient_checkpointing:
             model.gradient_checkpointing_enable()

        a  :    List[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 _lowercase(    self    )	->Tuple:
        a  :    str				=						self.prepare_config_and_inputs()

        (
            (
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							(
            a
        )			,							
        )  :    Any				=						config_and_inputs

        a  :    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 							_SCREAMING_SNAKE_CASE(   snake_case__       ,      snake_case__       ,      snake_case__       ,      unittest.TestCase	):
   UpperCamelCase_				:						Dict 			=   (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
   UpperCamelCase_				:						List[Any] 			=   (
       {'feature-extraction': RwkvModel, 'text-generation': RwkvForCausalLM} if is_torch_available() else {}
   )
   # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
   UpperCamelCase_				:						Optional[Any] 			=   False
   UpperCamelCase_				:						Dict 			=   False
   UpperCamelCase_				:						Optional[int] 			=   False
   UpperCamelCase_				:						int 			=   False
   UpperCamelCase_				:						Optional[int] 			=   False  # Rwkv does not support head masking
   def _lowercase(    self    )	->List[str]:
        a  :    Optional[int]				=						RwkvModelTester(self    )
        a  :    Dict				=						ConfigTester(
            self							,     config_class=lowerCAmelCase_							,     n_embd=37							,     common_properties=['''hidden_size''', '''num_hidden_layers''']    )
   def _lowercase(    self							,     lowerCAmelCase_							,     lowerCAmelCase_							,     lowerCAmelCase_=None    )	->Dict:
        if isinstance(lowerCAmelCase_							,     torch.Tensor    ):
             a			,							a  :    Union[str, Any]				=						member.max().item(), member.min().item()
        elif isinstance(lowerCAmelCase_							,     lowerCAmelCase_    ) or isinstance(lowerCAmelCase_							,     lowerCAmelCase_    ):
             a			,							a  :    Optional[Any]				=						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  :    Optional[int]				=						container

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

        if not is_inside_interval:
             a  :    Tuple				=						'''%s not found in %s''' % (safe_repr(lowerCAmelCase_    ), safe_repr(lowerCAmelCase_    ))
             self.fail(self._formatMessage(lowerCAmelCase_							,     lowerCAmelCase_    )    )
   def _lowercase(    self    )	->str:
        self.config_tester.run_common_tests()
   def _lowercase(    self    )	->Tuple:
        a  :    int				=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_rwkv_model(*lowerCAmelCase_    )
   def _lowercase(    self    )	->Dict:
        a  :    Dict				=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_causl_lm(*lowerCAmelCase_    )
   def _lowercase(    self    )	->List[str]:
        a  :    str				=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_state_equivalency(*lowerCAmelCase_    )
   def _lowercase(    self    )	->Optional[Any]:
        a			,							a  :    str				=						self.model_tester.prepare_config_and_inputs_for_common()

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

        a  :    Optional[int]				=						getattr(self.model_tester							,     '''seq_length'''							,     lowerCAmelCase_    )

        for model_class in self.all_model_classes:
             a  :    Tuple				=						True
             a  :    str				=						False
             a  :    str				=						True
             a  :    Union[str, Any]				=						model_class(lowerCAmelCase_    )
             model.to(lowerCAmelCase_    )
             model.eval()

             a  :    Optional[int]				=						self._prepare_for_class(lowerCAmelCase_							,     lowerCAmelCase_    )
             a  :    Dict				=						inputs['''input_ids'''].shape[0]
             with torch.no_grad():
                  a  :    int				=						model(**lowerCAmelCase_    )
             a  :    str				=						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  :    Tuple				=						self._prepare_for_class(lowerCAmelCase_							,     lowerCAmelCase_    )
             a  :    Tuple				=						inputs['''input_ids'''].shape[0]
             with torch.no_grad():
                  a  :    Any				=						model(**lowerCAmelCase_    )
             a  :    Union[str, Any]				=						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  :    Optional[Any]				=						len(lowerCAmelCase_    )

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

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

             a  :    List[str]				=						1
             self.assertEqual(out_len + added_hidden_states							,     len(lowerCAmelCase_    )    )

             a  :    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 _lowercase(    self    )	->str:
        for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             a  :    List[str]				=						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 							_SCREAMING_SNAKE_CASE(   unittest.TestCase	):
   def _lowercase(    self    )	->List[Any]:
        a  :    List[Any]				=						'''RWKV/rwkv-4-169m-pile'''
        a  :    Union[str, Any]				=						AutoTokenizer.from_pretrained(self.model_id    )
   def _lowercase(    self    )	->Tuple:
        a  :    Tuple				=						'''Hello my name is Jasmine and I am a newbie to the'''
        a  :    List[Any]				=						RwkvForCausalLM.from_pretrained(self.model_id    ).to(lowerCAmelCase_    )

        a  :    Optional[int]				=						self.tokenizer('''Hello my name is'''							,     return_tensors='''pt'''    ).input_ids.to(lowerCAmelCase_    )
        a  :    Optional[int]				=						model.generate(lowerCAmelCase_							,     max_new_tokens=10    )
        a  :    Dict				=						self.tokenizer.decode(output[0].tolist()    )

        self.assertEqual(lowerCAmelCase_							,     lowerCAmelCase_    )


   def _lowercase(    self    )	->Union[str, Any]:
        a  :    Any				=						'''Hello my name is Jasmine and I am a newbie to the'''

        a  :    str				=						self.tokenizer('''Hello my name is'''							,     return_tensors='''pt'''    ).input_ids.to(lowerCAmelCase_    )
        a  :    Any				=						RwkvForCausalLM.from_pretrained(self.model_id							,     torch_dtype=torch.bfloataa    ).to(lowerCAmelCase_    )

        a  :    Optional[Any]				=						model.generate(lowerCAmelCase_							,     max_new_tokens=10    )
        a  :    Optional[Any]				=						self.tokenizer.decode(output[0].tolist()    )

        self.assertEqual(lowerCAmelCase_							,     lowerCAmelCase_    )