


'''simple docstring'''





import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


__lowerCamelCase       				=	{
    '''vit_b32''': {
        '''embed_dim''': 512,
        '''image_resolution''': 768,
        '''context_length''': 16,
        '''vocab_size''': 4_9408,
        '''vision_layers''': 12,
        '''vision_width''': 768,
        '''vision_patch_size''': 32,
        '''transformer_width''': 512,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_b16''': {
        '''embed_dim''': 512,
        '''image_resolution''': 768,
        '''context_length''': 16,
        '''vocab_size''': 4_9408,
        '''vision_layers''': 12,
        '''vision_width''': 768,
        '''vision_patch_size''': 16,
        '''transformer_width''': 512,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_l14''': {
        '''embed_dim''': 768,
        '''image_resolution''': 840,
        '''context_length''': 16,
        '''vocab_size''': 4_9408,
        '''vision_layers''': 24,
        '''vision_width''': 1024,
        '''vision_patch_size''': 14,
        '''transformer_width''': 768,
        '''transformer_heads''': 12,
        '''transformer_layers''': 12,
    },
}


def 					a__    (				UpperCamelCase_	:			List[str],      UpperCamelCase_	:			Union[str, Any]="",      UpperCamelCase_	:			Any="/"	):
   UpperCAmelCase__  :List[str]            =			[]

   for k, v in params.items():
      UpperCAmelCase__  :Union[str, Any]            =			parent_key + sep + k if parent_key else k

      if isinstance(UpperCamelCase_,      collections.MutableMapping	):
         items.extend(flatten_nested_dict(UpperCamelCase_,      UpperCamelCase_,      sep=UpperCamelCase_	).items()	)
      else:
         items.append((new_key, v)	)
   return dict(UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			Optional[Any]	):
   return jax.tree_util.tree_map(lambda UpperCamelCase_	: x.astype(jnp.floataa	) if x.dtype == jnp.bfloataa else x,      UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			Optional[Any],      UpperCamelCase_	:			str	):
   UpperCAmelCase__      ,	UpperCAmelCase__      ,	UpperCAmelCase__  :List[Any]            =			pt_attn_layer.in_proj_weight.chunk(3,      dim=0	)
   UpperCAmelCase__      ,	UpperCAmelCase__      ,	UpperCAmelCase__  :List[str]            =			pt_attn_layer.in_proj_bias.chunk(3,      dim=0	)

   UpperCAmelCase__  :int            =			pt_attn_layer.out_proj.weight
   UpperCAmelCase__  :List[str]            =			pt_attn_layer.out_proj.bias

   UpperCAmelCase__  :Dict            =			q_proj
   UpperCAmelCase__  :List[str]            =			q_proj_bias

   UpperCAmelCase__  :Optional[Any]            =			k_proj
   UpperCAmelCase__  :Dict            =			k_proj_bias

   UpperCAmelCase__  :List[Any]            =			v_proj
   UpperCAmelCase__  :Dict            =			v_proj_bias

   UpperCAmelCase__  :Any            =			out_proj_weights
   UpperCAmelCase__  :int            =			out_proj_bias


def 					a__    (				UpperCamelCase_	:			Tuple,      UpperCamelCase_	:			Optional[int]	):
   copy_linear(hf_mlp.fca,      pt_mlp.c_fc	)
   copy_linear(hf_mlp.fca,      pt_mlp.c_proj	)


def 					a__    (				UpperCamelCase_	:			int,      UpperCamelCase_	:			List[Any]	):
   UpperCAmelCase__  :List[Any]            =			pt_linear.weight
   UpperCAmelCase__  :Dict            =			pt_linear.bias


def 					a__    (				UpperCamelCase_	:			Optional[Any],      UpperCamelCase_	:			List[str]	):
   # copy layer norms
   copy_linear(hf_layer.layer_norma,      pt_layer.ln_a	)
   copy_linear(hf_layer.layer_norma,      pt_layer.ln_a	)

   # copy MLP
   copy_mlp(hf_layer.mlp,      pt_layer.mlp	)

   # copy attn
   copy_attn_layer(hf_layer.self_attn,      pt_layer.attn	)


def 					a__    (				UpperCamelCase_	:			List[Any],      UpperCamelCase_	:			Tuple	):
   for hf_layer, pt_layer in zip(UpperCamelCase_,      UpperCamelCase_	):
      copy_layer(UpperCamelCase_,      UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			int,      UpperCamelCase_	:			Optional[Any]	):
   # copy  embeds
   UpperCAmelCase__  :List[Any]            =			pt_model.token_embedding.weight
   UpperCAmelCase__  :Any            =			pt_model.positional_embedding

   # copy layer norm
   copy_linear(hf_encoder.final_layer_norm,      pt_model.ln_final	)

   # copy hidden layers
   copy_layers(hf_encoder.encoder.layers,      pt_model.transformer.resblocks	)


def 					a__    (				UpperCamelCase_	:			str,      UpperCamelCase_	:			Optional[int]	):
   # copy projection
   UpperCAmelCase__  :Any            =			pt_model.text_projection.data.T

   # copy text encoder
   copy_encoder(hf_model.text_model,      UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			Optional[Any],      UpperCamelCase_	:			List[Any]	):
   # copy projection
   UpperCAmelCase__  :List[str]            =			pt_model.visual.proj.data.T

   # copy layer norms
   copy_linear(hf_model.vision_model.pre_layernorm,      pt_model.visual.ln_pre	)
   copy_linear(hf_model.vision_model.post_layernorm,      pt_model.visual.ln_post	)

   # copy embeds
   UpperCAmelCase__  :int            =			pt_model.visual.conva.weight.data
   UpperCAmelCase__  :Union[str, Any]            =			pt_model.visual.class_embedding
   UpperCAmelCase__  :Dict            =			pt_model.visual.positional_embedding.data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers,      pt_model.visual.transformer.resblocks	)


def 					a__    (				UpperCamelCase_	:			Union[str, Any],      UpperCamelCase_	:			Tuple	):
   UpperCAmelCase__  :int            =			flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']	)

   UpperCAmelCase__  :int            =			torch.from_numpy(flax_class_token_params['''scale''']	)
   UpperCAmelCase__  :Tuple            =			torch.from_numpy(flax_class_token_params['''bias''']	)
   UpperCAmelCase__  :Optional[int]            =			nn.Parameter(UpperCamelCase_	)
   UpperCAmelCase__  :List[Any]            =			nn.Parameter(UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			int,      UpperCamelCase_	:			List[Any]	):
   UpperCAmelCase__  :Tuple            =			hf_model.state_dict()
   UpperCAmelCase__  :Any            =			{}

   # Rename class prediction head flax params to pytorch HF
   UpperCAmelCase__  :Tuple            =			flatten_nested_dict(flax_params['''class_head''']	)

   for flax_key, v in flax_class_params.items():
      UpperCAmelCase__  :str            =			flax_key.replace('''/''',      '''.'''	)
      UpperCAmelCase__  :Dict            =			torch_key.replace('''.kernel''',      '''.weight'''	)
      UpperCAmelCase__  :str            =			torch_key.replace('''Dense_0''',      '''dense0'''	)
      UpperCAmelCase__  :Dict            =			'''class_head.''' + torch_key

      if "weight" in torch_key and v.ndim == 2:
         UpperCAmelCase__  :List[Any]            =			v.T

      UpperCAmelCase__  :int            =			nn.Parameter(torch.from_numpy(UpperCamelCase_	)	)

   # Rename box prediction box flax params to pytorch HF
   UpperCAmelCase__  :List[Any]            =			flatten_nested_dict(flax_params['''obj_box_head''']	)

   for flax_key, v in flax_box_params.items():
      UpperCAmelCase__  :Union[str, Any]            =			flax_key.replace('''/''',      '''.'''	)
      UpperCAmelCase__  :int            =			torch_key.replace('''.kernel''',      '''.weight'''	)
      UpperCAmelCase__  :Optional[int]            =			torch_key.replace('''_''',      ''''''	).lower()
      UpperCAmelCase__  :Optional[Any]            =			'''box_head.''' + torch_key

      if "weight" in torch_key and v.ndim == 2:
         UpperCAmelCase__  :Optional[int]            =			v.T

      UpperCAmelCase__  :Any            =			nn.Parameter(torch.from_numpy(UpperCamelCase_	)	)

   # Copy flax params to PyTorch params
   for name, param in new_params.items():
      if name in pt_params.keys():
         pt_params[name].copy_(UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			Union[str, Any],      UpperCamelCase_	:			Optional[Any]	):
   for k, v in flax_attn_params.items():
      if k.startswith('''transformer'''	):
         UpperCAmelCase__  :Union[str, Any]            =			k.replace('''transformer.resblocks''',      '''text_model.encoder.layers'''	)
      else:
         UpperCAmelCase__  :str            =			k.replace('''visual.transformer.resblocks''',      '''vision_model.encoder.layers'''	)

      UpperCAmelCase__  :List[Any]            =			torch_key.replace('''attn''',      '''self_attn'''	)
      UpperCAmelCase__  :int            =			torch_key.replace('''key''',      '''k_proj'''	)
      UpperCAmelCase__  :Any            =			torch_key.replace('''value''',      '''v_proj'''	)
      UpperCAmelCase__  :Optional[int]            =			torch_key.replace('''query''',      '''q_proj'''	)
      UpperCAmelCase__  :List[Any]            =			torch_key.replace('''out''',      '''out_proj'''	)

      if "bias" in torch_key and v.ndim == 2:
         UpperCAmelCase__  :int            =			v.shape[0] * v.shape[1]
         UpperCAmelCase__  :Tuple            =			v.reshape(UpperCamelCase_	)

      if "weight" in torch_key and "out" in torch_key:
         UpperCAmelCase__  :List[Any]            =			(v.shape[0] * v.shape[1], v.shape[2])
         UpperCAmelCase__  :Dict            =			v.reshape(UpperCamelCase_	).T

      if "weight" in torch_key and "out" not in torch_key:
         UpperCAmelCase__  :str            =			(v.shape[0], v.shape[1] * v.shape[2])
         UpperCAmelCase__  :Dict            =			v.reshape(UpperCamelCase_	).T

      # Copy flax CLIP attn params to HF PyTorch params
      UpperCAmelCase__  :Any            =			torch.from_numpy(UpperCamelCase_	)
      hf_backbone.state_dict()[torch_key].copy_(UpperCamelCase_	)


def 					a__    (				UpperCamelCase_	:			List[str]	):
   UpperCAmelCase__  :str            =			{}
   UpperCAmelCase__  :Union[str, Any]            =			[]

   for k, v in params.items():
      if "attn." in k:
         UpperCAmelCase__  :List[Any]            =			k[: k.rindex('''attn.'''	) + 5]
         if base in processed_attn_layers:
            continue

         processed_attn_layers.append(UpperCamelCase_	)
         UpperCAmelCase__  :Any            =			params[base + '''out.weight'''].shape[-1]
         UpperCAmelCase__  :Optional[Any]            =			params[base + '''out.weight'''].reshape(UpperCamelCase_,      UpperCamelCase_	).T
         UpperCAmelCase__  :Tuple            =			params[base + '''out.bias''']
      else:
         UpperCAmelCase__  :List[str]            =			v
   return new_params


def 					a__    (				UpperCamelCase_	:			Union[str, Any],      UpperCamelCase_	:			int	):
   UpperCAmelCase__  :Optional[Any]            =			CLIP(**UpperCamelCase_	)
   torch_model.eval()
   UpperCAmelCase__  :List[str]            =			torch_model.state_dict()

   UpperCAmelCase__  :str            =			flatten_nested_dict(flax_params['''backbone''']['''clip''']	)
   UpperCAmelCase__  :List[str]            =			{}

   for flax_key, v in flax_clip_params.items():
      UpperCAmelCase__  :List[str]            =			flax_key.replace('''/''',      '''.'''	)
      UpperCAmelCase__  :str            =			torch_key.replace('''text.token_embedding.embedding''',      '''token_embedding.kernel'''	)

      if (
          torch_key.startswith('''text.transformer'''	)
          or torch_key.startswith('''text.text_projection'''	)
          or torch_key.startswith('''text.ln_final'''	)
          or torch_key.startswith('''text.positional_embedding'''	)
      ):
         UpperCAmelCase__  :Dict            =			torch_key[5:]

      UpperCAmelCase__  :List[Any]            =			torch_key.replace('''text_projection.kernel''',      '''text_projection'''	)
      UpperCAmelCase__  :List[str]            =			torch_key.replace('''visual.proj.kernel''',      '''visual.proj'''	)
      UpperCAmelCase__  :int            =			torch_key.replace('''.scale''',      '''.weight'''	)
      UpperCAmelCase__  :List[str]            =			torch_key.replace('''.kernel''',      '''.weight'''	)

      if "conv" in torch_key or "downsample.0.weight" in torch_key:
         UpperCAmelCase__  :Optional[Any]            =			v.transpose(3,      2,      0,      1	)

      elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
         # Fully connected layers are transposed, embeddings are not
         UpperCAmelCase__  :Union[str, Any]            =			v.T

      UpperCAmelCase__  :Optional[Any]            =			v

   UpperCAmelCase__  :Dict            =			_convert_attn_layers(UpperCamelCase_	)
   new_torch_params.update(UpperCamelCase_	)
   UpperCAmelCase__  :List[str]            =			{}

   # Copy flax CLIP backbone params to PyTorch params
   for name, param in new_torch_params.items():
      if name in torch_clip_params.keys():
         UpperCAmelCase__  :int            =			torch.from_numpy(new_torch_params[name]	)
         torch_clip_params[name].copy_(UpperCamelCase_	)
      else:
         UpperCAmelCase__  :Union[str, Any]            =			param

   return torch_clip_params, torch_model, attn_params


@torch.no_grad()
def 					a__    (				UpperCamelCase_	:			Optional[int],      UpperCamelCase_	:			Any,      UpperCamelCase_	:			List[Any],      UpperCamelCase_	:			Optional[Any],      UpperCamelCase_	:			Dict=None	):
   UpperCAmelCase__  :List[Any]            =			Repository(UpperCamelCase_,      clone_from=F'''google/{pytorch_dump_folder_path}'''	)
   repo.git_pull()

   if config_path is not None:
      UpperCAmelCase__  :List[Any]            =			OwlViTConfig.from_pretrained(UpperCamelCase_	)
   else:
      UpperCAmelCase__  :Union[str, Any]            =			OwlViTConfig()

   UpperCAmelCase__  :Dict            =			OwlViTModel(UpperCamelCase_	).eval()
   UpperCAmelCase__  :Any            =			OwlViTForObjectDetection(UpperCamelCase_	).eval()

   copy_text_model_and_projection(UpperCamelCase_,      UpperCamelCase_	)
   copy_vision_model_and_projection(UpperCamelCase_,      UpperCamelCase_	)
   UpperCAmelCase__  :Optional[int]            =			pt_backbone.logit_scale
   copy_flax_attn_params(UpperCamelCase_,      UpperCamelCase_	)

   UpperCAmelCase__  :List[str]            =			hf_backbone
   copy_class_merge_token(UpperCamelCase_,      UpperCamelCase_	)
   copy_class_box_heads(UpperCamelCase_,      UpperCamelCase_	)

   # Save HF model
   hf_model.save_pretrained(repo.local_dir	)

   # Initialize image processor
   UpperCAmelCase__  :Any            =			OwlViTImageProcessor(
       size=config.vision_config.image_size,      crop_size=config.vision_config.image_size	)
   # Initialize tokenizer
   UpperCAmelCase__  :int            =			CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32''',      pad_token='''!''',      model_max_length=16	)

   # Initialize processor
   UpperCAmelCase__  :Optional[int]            =			OwlViTProcessor(image_processor=UpperCamelCase_,      tokenizer=UpperCamelCase_	)
   image_processor.save_pretrained(repo.local_dir	)
   processor.save_pretrained(repo.local_dir	)

   repo.git_add()
   repo.git_commit('''Upload model and processor'''	)
   repo.git_push()


if __name__ == "__main__":
    __lowerCamelCase       				=	argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--owlvit_version''',
        default=None,
        type=str,
        required=True,
        help='''OWL-ViT model name [clip_b16, clip_b32, clip_l14].''',
    )
    parser.add_argument(
        '''--owlvit_checkpoint''', default=None, type=str, required=True, help='''Path to flax model checkpoint.'''
    )
    parser.add_argument('''--hf_config''', default=None, type=str, required=True, help='''Path to HF model config.''')
    parser.add_argument(
        '''--pytorch_dump_folder_path''', default='''hf_model''', type=str, help='''Path to the output PyTorch model.'''
    )
    __lowerCamelCase       				=	parser.parse_args()

    # Initialize PyToch clip model
    __lowerCamelCase       				=	args.owlvit_version
    if model_name == "clip_b16":
        __lowerCamelCase       				=	CONFIGS['''vit_b16''']
    elif model_name == "clip_b32":
        __lowerCamelCase       				=	CONFIGS['''vit_b32''']
    elif model_name == "clip_l14":
        __lowerCamelCase       				=	CONFIGS['''vit_l14''']

    # Load from checkpoint and convert params to float-32
    __lowerCamelCase       				=	checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['''optimizer''']['''target''']
    __lowerCamelCase       				=	jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
    del variables

    # Convert CLIP backbone
    __lowerCamelCase				,		__lowerCamelCase				,		__lowerCamelCase       				=	convert_clip_backbone(flax_params, torch_config)

    convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)







'''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:
    __lowerCamelCase       				=	False




class 							UpperCAmelCase :


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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Tuple ,					__lowerCamelCase       :  Dict=False ,					__lowerCamelCase       :  Dict=False ,					__lowerCamelCase       :  List[str]=False   ):
    UpperCAmelCase__  :Union[str, Any]            =			ids_tensor([self.batch_size, self.seq_length] ,					self.vocab_size   )

    UpperCAmelCase__  :Dict            =			None
    if self.use_input_mask:
       UpperCAmelCase__  :Optional[int]            =			random_attention_mask([self.batch_size, self.seq_length]   )

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

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

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

    UpperCAmelCase__  :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 					__SCREAMING_SNAKE_CASE				(			self       :  str ,					__lowerCamelCase       :  Union[str, Any]=False ,					__lowerCamelCase       :  int=False ,					__lowerCamelCase       :  int=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 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    UpperCAmelCase__  :int            =			self.get_config()
    UpperCAmelCase__  :Any            =			3_0_0
    return config


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Tuple   ):
    (
        (
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	
    )  :Dict            =			self.prepare_config_and_inputs()

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

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

    UpperCAmelCase__  :int            =			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 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int] ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  Tuple ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  Union[str, Any] ,					__lowerCamelCase       :  Tuple ,					*__lowerCamelCase       :  Optional[Any]   ):
    UpperCAmelCase__  :Optional[Any]            =			RwkvModel(config=__lowerCamelCase   )
    model.to(__lowerCamelCase   )
    model.eval()

    UpperCAmelCase__  :Tuple            =			model(__lowerCamelCase   )
    UpperCAmelCase__  :int            =			outputs.last_hidden_state

    UpperCAmelCase__  :Tuple            =			model(input_ids[:, :2]   )
    UpperCAmelCase__  :List[Any]            =			outputs.last_hidden_state

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

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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str ,					__lowerCamelCase       :  int ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  List[Any] ,					__lowerCamelCase       :  List[str] ,					__lowerCamelCase       :  Any ,					*__lowerCamelCase       :  Any ,					__lowerCamelCase       :  Tuple=False   ):
    UpperCAmelCase__  :Any            =			RwkvForCausalLM(__lowerCamelCase   )
    model.to(__lowerCamelCase   )
    if gradient_checkpointing:
       model.gradient_checkpointing_enable()

    UpperCAmelCase__  :Union[str, 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 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__  :Optional[Any]            =			self.prepare_config_and_inputs()

    (
        (
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	
    )  :Any            =			config_and_inputs

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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    UpperCAmelCase__  :Dict            =			RwkvModelTester(self   )
    UpperCAmelCase__  :Union[str, Any]            =			ConfigTester(
        self ,					config_class=__lowerCamelCase ,					n_embd=3_7 ,					common_properties=['''hidden_size''', '''num_hidden_layers''']   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict ,					__lowerCamelCase       :  Optional[Any] ,					__lowerCamelCase       :  int ,					__lowerCamelCase       :  Optional[int]=None   ):
    if isinstance(__lowerCamelCase ,					torch.Tensor   ):
       UpperCAmelCase__      ,	UpperCAmelCase__  :Dict            =			member.max().item(), member.min().item()
    elif isinstance(__lowerCamelCase ,					__lowerCamelCase   ) or isinstance(__lowerCamelCase ,					__lowerCamelCase   ):
       UpperCAmelCase__      ,	UpperCAmelCase__  :List[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'''   )

    UpperCAmelCase__      ,	UpperCAmelCase__  :Any            =			container

    UpperCAmelCase__  :List[str]            =			(min_value >= expected_min) and (max_value <= expected_max)

    if not is_inside_interval:
       UpperCAmelCase__  :int            =			'''%s not found in %s''' % (safe_repr(__lowerCamelCase   ), safe_repr(__lowerCamelCase   ))
       self.fail(self._formatMessage(__lowerCamelCase ,					__lowerCamelCase   )   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    self.config_tester.run_common_tests()


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    UpperCAmelCase__  :Tuple            =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_rwkv_model(*__lowerCamelCase   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  int   ):
    UpperCAmelCase__  :str            =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_causl_lm(*__lowerCamelCase   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    UpperCAmelCase__  :Union[str, Any]            =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_state_equivalency(*__lowerCamelCase   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    UpperCAmelCase__      ,	UpperCAmelCase__  :Dict            =			self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
       UpperCAmelCase__  :Optional[int]            =			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 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__      ,	UpperCAmelCase__  :int            =			self.model_tester.prepare_config_and_inputs_for_common()
    UpperCAmelCase__  :Dict            =			True

    UpperCAmelCase__  :Dict            =			getattr(self.model_tester ,					'''seq_length''' ,					__lowerCamelCase   )

    for model_class in self.all_model_classes:
       UpperCAmelCase__  :Tuple            =			True
       UpperCAmelCase__  :Optional[int]            =			False
       UpperCAmelCase__  :List[Any]            =			True
       UpperCAmelCase__  :Optional[int]            =			model_class(__lowerCamelCase   )
       model.to(__lowerCamelCase   )
       model.eval()

       UpperCAmelCase__  :List[Any]            =			self._prepare_for_class(__lowerCamelCase ,					__lowerCamelCase   )
       UpperCAmelCase__  :Tuple            =			inputs['''input_ids'''].shape[0]
       with torch.no_grad():
          UpperCAmelCase__  :List[str]            =			model(**__lowerCamelCase   )
       UpperCAmelCase__  :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"]
       UpperCAmelCase__  :Optional[Any]            =			True
       UpperCAmelCase__  :Union[str, Any]            =			model_class(__lowerCamelCase   )
       model.to(__lowerCamelCase   )
       model.eval()

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

       # Check attention is always last and order is fine
       UpperCAmelCase__  :int            =			True
       UpperCAmelCase__  :Tuple            =			True
       UpperCAmelCase__  :List[Any]            =			model_class(__lowerCamelCase   )
       model.to(__lowerCamelCase   )
       model.eval()

       UpperCAmelCase__  :Optional[int]            =			self._prepare_for_class(__lowerCamelCase ,					__lowerCamelCase   )
       UpperCAmelCase__  :int            =			inputs['''input_ids'''].shape[0]
       with torch.no_grad():
          UpperCAmelCase__  :Optional[int]            =			model(**__lowerCamelCase   )

       UpperCAmelCase__  :List[Any]            =			1
       self.assertEqual(out_len + added_hidden_states ,					len(__lowerCamelCase   )   )

       UpperCAmelCase__  :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 					__SCREAMING_SNAKE_CASE				(			self       :  Dict   ):
    for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       UpperCAmelCase__  :List[Any]            =			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 							UpperCAmelCase ( unittest.TestCase	):


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    UpperCAmelCase__  :List[Any]            =			'''RWKV/rwkv-4-169m-pile'''
    UpperCAmelCase__  :str            =			AutoTokenizer.from_pretrained(self.model_id   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Any   ):
    UpperCAmelCase__  :Optional[int]            =			'''Hello my name is Jasmine and I am a newbie to the'''
    UpperCAmelCase__  :Union[str, Any]            =			RwkvForCausalLM.from_pretrained(self.model_id   ).to(__lowerCamelCase   )

    UpperCAmelCase__  :List[str]            =			self.tokenizer('''Hello my name is''' ,					return_tensors='''pt'''   ).input_ids.to(__lowerCamelCase   )
    UpperCAmelCase__  :str            =			model.generate(__lowerCamelCase ,					max_new_tokens=1_0   )
    UpperCAmelCase__  :str            =			self.tokenizer.decode(output[0].tolist()   )

    self.assertEqual(__lowerCamelCase ,					__lowerCamelCase   )





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

    UpperCAmelCase__  :int            =			self.tokenizer('''Hello my name is''' ,					return_tensors='''pt'''   ).input_ids.to(__lowerCamelCase   )
    UpperCAmelCase__  :Dict            =			RwkvForCausalLM.from_pretrained(self.model_id ,					torch_dtype=torch.bfloataa   ).to(__lowerCamelCase   )

    UpperCAmelCase__  :str            =			model.generate(__lowerCamelCase ,					max_new_tokens=1_0   )
    UpperCAmelCase__  :int            =			self.tokenizer.decode(output[0].tolist()   )

    self.assertEqual(__lowerCamelCase ,					__lowerCamelCase   )



