

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							:  List[Any]     =				{
    "vit_b32": {
        "embed_dim": 5_1_2,
        "image_resolution": 7_6_8,
        "context_length": 1_6,
        "vocab_size": 4_9_4_0_8,
        "vision_layers": 1_2,
        "vision_width": 7_6_8,
        "vision_patch_size": 3_2,
        "transformer_width": 5_1_2,
        "transformer_heads": 8,
        "transformer_layers": 1_2,
    },
    "vit_b16": {
        "embed_dim": 5_1_2,
        "image_resolution": 7_6_8,
        "context_length": 1_6,
        "vocab_size": 4_9_4_0_8,
        "vision_layers": 1_2,
        "vision_width": 7_6_8,
        "vision_patch_size": 1_6,
        "transformer_width": 5_1_2,
        "transformer_heads": 8,
        "transformer_layers": 1_2,
    },
    "vit_l14": {
        "embed_dim": 7_6_8,
        "image_resolution": 8_4_0,
        "context_length": 1_6,
        "vocab_size": 4_9_4_0_8,
        "vision_layers": 2_4,
        "vision_width": 1_0_2_4,
        "vision_patch_size": 1_4,
        "transformer_width": 7_6_8,
        "transformer_heads": 1_2,
        "transformer_layers": 1_2,
    },
}

def      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     Dict=""							,					UpperCAmelCase				:     Optional[Any]="/"     )    ->		Optional[int]:
  UpperCAmelCase			:				str										=					[]

  for k, v in params.items():
    UpperCAmelCase			:				int										=					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[int]     )    ->		Tuple:
  return jax.tree_util.tree_map(lambda UpperCAmelCase     : x.astype(jnp.floataa     ) if x.dtype == jnp.bfloataa else x							,					UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     Union[str, Any]							,					UpperCAmelCase				:     List[Any]     )    ->		str:
  UpperCAmelCase					,		UpperCAmelCase					,		UpperCAmelCase			:				List[Any]										=					pt_attn_layer.in_proj_weight.chunk(3							,					dim=0     )
  UpperCAmelCase					,		UpperCAmelCase					,		UpperCAmelCase			:				Tuple										=					pt_attn_layer.in_proj_bias.chunk(3							,					dim=0     )

  UpperCAmelCase			:				Optional[Any]										=					pt_attn_layer.out_proj.weight
  UpperCAmelCase			:				Tuple										=					pt_attn_layer.out_proj.bias

  UpperCAmelCase			:				List[str]										=					q_proj
  UpperCAmelCase			:				Optional[int]										=					q_proj_bias

  UpperCAmelCase			:				Any										=					k_proj
  UpperCAmelCase			:				Optional[Any]										=					k_proj_bias

  UpperCAmelCase			:				Optional[Any]										=					v_proj
  UpperCAmelCase			:				Optional[Any]										=					v_proj_bias

  UpperCAmelCase			:				str										=					out_proj_weights
  UpperCAmelCase			:				List[str]										=					out_proj_bias

def      a__	(						UpperCAmelCase				:     Any							,					UpperCAmelCase				:     Dict     )    ->		Optional[int]:
  copy_linear(hf_mlp.fca							,					pt_mlp.c_fc     )
  copy_linear(hf_mlp.fca							,					pt_mlp.c_proj     )

def      a__	(						UpperCAmelCase				:     Optional[int]							,					UpperCAmelCase				:     str     )    ->		Dict:
  UpperCAmelCase			:				Dict										=					pt_linear.weight
  UpperCAmelCase			:				List[Any]										=					pt_linear.bias

def      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     int     )    ->		Tuple:
  # 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				:     Tuple							,					UpperCAmelCase				:     Union[str, Any]     )    ->		int:
  for hf_layer, pt_layer in zip(UpperCAmelCase							,					UpperCAmelCase     ):
    copy_layer(UpperCAmelCase							,					UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     int							,					UpperCAmelCase				:     Dict     )    ->		Union[str, 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				:     int							,					UpperCAmelCase				:     Tuple     )    ->		str:
  # copy projection
  UpperCAmelCase			:				int										=					pt_model.text_projection.data.T

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

def      a__	(						UpperCAmelCase				:     int							,					UpperCAmelCase				:     List[str]     )    ->		List[Any]:
  # copy projection
  UpperCAmelCase			:				List[Any]										=					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			:				Dict										=					pt_model.visual.conva.weight.data
  UpperCAmelCase			:				Any										=					pt_model.visual.class_embedding
  UpperCAmelCase			:				str										=					pt_model.visual.positional_embedding.data

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

def      a__	(						UpperCAmelCase				:     Any							,					UpperCAmelCase				:     Any     )    ->		Union[str, Any]:
  UpperCAmelCase			:				str										=					flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']     )

  UpperCAmelCase			:				str										=					torch.from_numpy(flax_class_token_params['''scale''']     )
  UpperCAmelCase			:				Union[str, Any]										=					torch.from_numpy(flax_class_token_params['''bias''']     )
  UpperCAmelCase			:				Union[str, Any]										=					nn.Parameter(UpperCAmelCase     )
  UpperCAmelCase			:				int										=					nn.Parameter(UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     Any							,					UpperCAmelCase				:     Optional[Any]     )    ->		List[Any]:
  UpperCAmelCase			:				int										=					hf_model.state_dict()
  UpperCAmelCase			:				Any										=					{}

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

  for flax_key, v in flax_class_params.items():
    UpperCAmelCase			:				List[str]										=					flax_key.replace('''/'''							,					'''.'''     )
    UpperCAmelCase			:				Dict										=					torch_key.replace('''.kernel'''							,					'''.weight'''     )
    UpperCAmelCase			:				Optional[int]										=					torch_key.replace('''Dense_0'''							,					'''dense0'''     )
    UpperCAmelCase			:				List[Any]										=					'''class_head.''' + torch_key

    if "weight" in torch_key and v.ndim == 2:
      UpperCAmelCase			:				Union[str, Any]										=					v.T

    UpperCAmelCase			:				List[str]										=					nn.Parameter(torch.from_numpy(UpperCAmelCase     )     )

  # Rename box prediction box flax params to pytorch HF
  UpperCAmelCase			:				int										=					flatten_nested_dict(flax_params['''obj_box_head''']     )

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

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

    UpperCAmelCase			:				int										=					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				:     int							,					UpperCAmelCase				:     List[Any]     )    ->		List[Any]:
  for k, v in flax_attn_params.items():
    if k.startswith('''transformer'''     ):
      UpperCAmelCase			:				Optional[int]										=					k.replace('''transformer.resblocks'''							,					'''text_model.encoder.layers'''     )
    else:
      UpperCAmelCase			:				Tuple										=					k.replace('''visual.transformer.resblocks'''							,					'''vision_model.encoder.layers'''     )

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

    if "bias" in torch_key and v.ndim == 2:
      UpperCAmelCase			:				str										=					v.shape[0] * v.shape[1]
      UpperCAmelCase			:				str										=					v.reshape(UpperCAmelCase     )

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

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

    # Copy flax CLIP attn params to HF PyTorch params
    UpperCAmelCase			:				List[str]										=					torch.from_numpy(UpperCAmelCase     )
    hf_backbone.state_dict()[torch_key].copy_(UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     Union[str, Any]     )    ->		List[str]:
  UpperCAmelCase			:				Optional[Any]										=					{}
  UpperCAmelCase			:				Dict										=					[]

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

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

def      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     Optional[Any]     )    ->		int:
  UpperCAmelCase			:				Optional[int]										=					CLIP(**UpperCAmelCase     )
  torch_model.eval()
  UpperCAmelCase			:				Optional[Any]										=					torch_model.state_dict()

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

  for flax_key, v in flax_clip_params.items():
    UpperCAmelCase			:				Optional[int]										=					flax_key.replace('''/'''							,					'''.'''     )
    UpperCAmelCase			:				Dict										=					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			:				List[str]										=					torch_key[5:]

    UpperCAmelCase			:				Dict										=					torch_key.replace('''text_projection.kernel'''							,					'''text_projection'''     )
    UpperCAmelCase			:				Tuple										=					torch_key.replace('''visual.proj.kernel'''							,					'''visual.proj'''     )
    UpperCAmelCase			:				Dict										=					torch_key.replace('''.scale'''							,					'''.weight'''     )
    UpperCAmelCase			:				str										=					torch_key.replace('''.kernel'''							,					'''.weight'''     )

    if "conv" in torch_key or "downsample.0.weight" in torch_key:
      UpperCAmelCase			:				Dict										=					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			:				Dict										=					v.T

    UpperCAmelCase			:				Union[str, Any]										=					v

  UpperCAmelCase			:				Dict										=					_convert_attn_layers(UpperCAmelCase     )
  new_torch_params.update(UpperCAmelCase     )
  UpperCAmelCase			:				int										=					{}

  # Copy flax CLIP backbone params to PyTorch params
  for name, param in new_torch_params.items():
    if name in torch_clip_params.keys():
      UpperCAmelCase			:				Optional[Any]										=					torch.from_numpy(new_torch_params[name]     )
      torch_clip_params[name].copy_(UpperCAmelCase     )
    else:
      UpperCAmelCase			:				List[Any]										=					param

  return torch_clip_params, torch_model, attn_params

@torch.no_grad()
def      a__	(						UpperCAmelCase				:     str							,					UpperCAmelCase				:     Union[str, Any]							,					UpperCAmelCase				:     Optional[int]							,					UpperCAmelCase				:     List[str]							,					UpperCAmelCase				:     Optional[Any]=None     )    ->		Optional[int]:
  UpperCAmelCase			:				Optional[int]										=					Repository(UpperCAmelCase							,					clone_from=f'''google/{pytorch_dump_folder_path}'''     )
  repo.git_pull()

  if config_path is not None:
    UpperCAmelCase			:				str										=					OwlViTConfig.from_pretrained(UpperCAmelCase     )
  else:
    UpperCAmelCase			:				List[str]										=					OwlViTConfig()

  UpperCAmelCase			:				List[str]										=					OwlViTModel(UpperCAmelCase     ).eval()
  UpperCAmelCase			:				Optional[int]										=					OwlViTForObjectDetection(UpperCAmelCase     ).eval()

  copy_text_model_and_projection(UpperCAmelCase							,					UpperCAmelCase     )
  copy_vision_model_and_projection(UpperCAmelCase							,					UpperCAmelCase     )
  UpperCAmelCase			:				Tuple										=					pt_backbone.logit_scale
  copy_flax_attn_params(UpperCAmelCase							,					UpperCAmelCase     )

  UpperCAmelCase			:				int										=					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			:				List[str]										=					OwlViTImageProcessor(
      size=config.vision_config.image_size							,					crop_size=config.vision_config.image_size     )
  # Initialize tokenizer
  UpperCAmelCase			:				Union[str, Any]										=					CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''							,					pad_token='''!'''							,					model_max_length=16     )

  # Initialize processor
  UpperCAmelCase			:				str										=					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							:  List[str]     =				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							:  List[Any]     =				parser.parse_args()

     # Initialize PyToch clip model
     _lowerCamelCase							:  Any     =				args.owlvit_version
     if model_name == "clip_b16":
          _lowerCamelCase							:  Optional[int]     =				CONFIGS["vit_b16"]
     elif model_name == "clip_b32":
          _lowerCamelCase							:  Dict     =				CONFIGS["vit_b32"]
     elif model_name == "clip_l14":
          _lowerCamelCase							:  Optional[Any]     =				CONFIGS["vit_l14"]

     # Load from checkpoint and convert params to float-32
     _lowerCamelCase							:  Optional[int]     =				checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)["optimizer"]["target"]
     _lowerCamelCase							:  List[str]     =				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							:  Optional[Any]     =				convert_clip_backbone(flax_params, torch_config)

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








import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         FunnelBaseModel,
         FunnelForMaskedLM,
         FunnelForMultipleChoice,
         FunnelForPreTraining,
         FunnelForQuestionAnswering,
         FunnelForSequenceClassification,
         FunnelForTokenClassification,
         FunnelModel,
     )





class    __UpperCAmelCase	:



   def __init__(				self				:	List[Any],						__A				:	str,						__A				:	Union[str, Any]=1_3,						__A				:	Union[str, Any]=7,						__A				:	List[str]=True,						__A				:	Optional[Any]=True,						__A				:	Tuple=True,						__A				:	Any=True,						__A				:	Dict=9_9,						__A				:	Union[str, Any]=[1, 1, 2],						__A				:	Tuple=1,						__A				:	int=3_2,						__A				:	List[str]=4,						__A				:	List[str]=8,						__A				:	int=3_7,						__A				:	Optional[int]="gelu_new",						__A				:	List[Any]=0.1,						__A				:	str=0.1,						__A				:	int=0.0,						__A				:	Union[str, Any]=5_1_2,						__A				:	int=3,						__A				:	str=0.0_2,						__A				:	Tuple=3,						__A				:	Any=4,						__A				:	Any=None,						__A				:	Any=False,						):
     UpperCAmelCase			:				str										=					parent
     UpperCAmelCase			:				Optional[Any]										=					batch_size
     UpperCAmelCase			:				str										=					seq_length
     UpperCAmelCase			:				str										=					is_training
     UpperCAmelCase			:				Tuple										=					use_input_mask
     UpperCAmelCase			:				Union[str, Any]										=					use_token_type_ids
     UpperCAmelCase			:				Dict										=					use_labels
     UpperCAmelCase			:				Union[str, Any]										=					vocab_size
     UpperCAmelCase			:				Tuple										=					block_sizes
     UpperCAmelCase			:				List[Any]										=					num_decoder_layers
     UpperCAmelCase			:				Union[str, Any]										=					d_model
     UpperCAmelCase			:				Tuple										=					n_head
     UpperCAmelCase			:				Dict										=					d_head
     UpperCAmelCase			:				Optional[int]										=					d_inner
     UpperCAmelCase			:				Optional[int]										=					hidden_act
     UpperCAmelCase			:				int										=					hidden_dropout
     UpperCAmelCase			:				Any										=					attention_dropout
     UpperCAmelCase			:				str										=					activation_dropout
     UpperCAmelCase			:				Dict										=					max_position_embeddings
     UpperCAmelCase			:				Optional[Any]										=					type_vocab_size
     UpperCAmelCase			:				Any										=					2
     UpperCAmelCase			:				List[Any]										=					num_labels
     UpperCAmelCase			:				str										=					num_choices
     UpperCAmelCase			:				str										=					scope
     UpperCAmelCase			:				int										=					initializer_std

     # Used in the tests to check the size of the first attention layer
     UpperCAmelCase			:				int										=					n_head
     # Used in the tests to check the size of the first hidden state
     UpperCAmelCase			:				Dict										=					self.d_model
     # Used in the tests to check the number of output hidden states/attentions
     UpperCAmelCase			:				List[str]										=					sum(self.block_sizes		) + (0 if base else self.num_decoder_layers)
     # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
     # the last hidden state of the first block (which is the first hidden state of the decoder).
     if not base:
       UpperCAmelCase			:				List[str]										=					self.num_hidden_layers + 2



   def 			__magic_name__       (				self				:	List[str]		):
     UpperCAmelCase			:				Dict										=					ids_tensor([self.batch_size, self.seq_length],						self.vocab_size		)

     UpperCAmelCase			:				Optional[int]										=					None
     if self.use_input_mask:
       UpperCAmelCase			:				List[str]										=					ids_tensor([self.batch_size, self.seq_length],						vocab_size=2		)

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

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

     UpperCAmelCase			:				Dict										=					self.get_config()

     return (
         config,
         input_ids,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
         choice_labels,
         fake_token_labels,
     )



   def 			__magic_name__       (				self				:	List[Any]		):
     return FunnelConfig(
         vocab_size=self.vocab_size,						block_sizes=self.block_sizes,						num_decoder_layers=self.num_decoder_layers,						d_model=self.d_model,						n_head=self.n_head,						d_head=self.d_head,						d_inner=self.d_inner,						hidden_act=self.hidden_act,						hidden_dropout=self.hidden_dropout,						attention_dropout=self.attention_dropout,						activation_dropout=self.activation_dropout,						max_position_embeddings=self.max_position_embeddings,						type_vocab_size=self.type_vocab_size,						initializer_std=self.initializer_std,						)



   def 			__magic_name__       (				self				:	Union[str, Any],						__A				:	List[str],						__A				:	Dict,						__A				:	Optional[int],						__A				:	List[Any],						__A				:	List[Any],						__A				:	Optional[int],						__A				:	Dict,						__A				:	Dict,						):
     UpperCAmelCase			:				List[str]										=					FunnelModel(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				List[Any]										=					model(__A,						attention_mask=__A,						token_type_ids=__A		)
     UpperCAmelCase			:				Optional[int]										=					model(__A,						token_type_ids=__A		)
     UpperCAmelCase			:				Optional[int]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, self.seq_length, self.d_model)		)

     UpperCAmelCase			:				Optional[int]										=					False
     UpperCAmelCase			:				List[str]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, self.seq_length, self.d_model)		)

     UpperCAmelCase			:				List[str]										=					False
     UpperCAmelCase			:				Dict										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, self.seq_length, self.d_model)		)



   def 			__magic_name__       (				self				:	Dict,						__A				:	str,						__A				:	List[Any],						__A				:	Union[str, Any],						__A				:	Union[str, Any],						__A				:	Tuple,						__A				:	Tuple,						__A				:	List[Any],						__A				:	Dict,						):
     UpperCAmelCase			:				List[Any]										=					FunnelBaseModel(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				int										=					model(__A,						attention_mask=__A,						token_type_ids=__A		)
     UpperCAmelCase			:				List[Any]										=					model(__A,						token_type_ids=__A		)
     UpperCAmelCase			:				Optional[Any]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, 2, self.d_model)		)

     UpperCAmelCase			:				str										=					False
     UpperCAmelCase			:				Optional[Any]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, 3, self.d_model)		)

     UpperCAmelCase			:				Tuple										=					False
     UpperCAmelCase			:				Optional[Any]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, 2, self.d_model)		)



   def 			__magic_name__       (				self				:	Optional[int],						__A				:	Tuple,						__A				:	Dict,						__A				:	str,						__A				:	Dict,						__A				:	Union[str, Any],						__A				:	List[Any],						__A				:	Any,						__A				:	str,						):
     UpperCAmelCase			:				int										=					self.num_labels
     UpperCAmelCase			:				int										=					FunnelForPreTraining(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				List[Any]										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.seq_length)		)



   def 			__magic_name__       (				self				:	List[str],						__A				:	Optional[int],						__A				:	Union[str, Any],						__A				:	int,						__A				:	Union[str, Any],						__A				:	Optional[int],						__A				:	str,						__A				:	int,						__A				:	str,						):
     UpperCAmelCase			:				List[Any]										=					FunnelForMaskedLM(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Optional[int]										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.seq_length, self.vocab_size)		)



   def 			__magic_name__       (				self				:	Optional[int],						__A				:	Optional[Any],						__A				:	Optional[Any],						__A				:	Dict,						__A				:	Any,						__A				:	Optional[Any],						__A				:	Tuple,						__A				:	Union[str, Any],						__A				:	List[str],						):
     UpperCAmelCase			:				str										=					self.num_labels
     UpperCAmelCase			:				Optional[int]										=					FunnelForSequenceClassification(__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Dict										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.num_labels)		)



   def 			__magic_name__       (				self				:	Union[str, Any],						__A				:	Any,						__A				:	List[str],						__A				:	str,						__A				:	List[Any],						__A				:	List[str],						__A				:	Optional[Any],						__A				:	List[Any],						__A				:	List[str],						):
     UpperCAmelCase			:				str										=					self.num_choices
     UpperCAmelCase			:				Optional[int]										=					FunnelForMultipleChoice(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				List[Any]										=					input_ids.unsqueeze(1		).expand(-1,						self.num_choices,						-1		).contiguous()
     UpperCAmelCase			:				List[Any]										=					token_type_ids.unsqueeze(1		).expand(-1,						self.num_choices,						-1		).contiguous()
     UpperCAmelCase			:				List[Any]										=					input_mask.unsqueeze(1		).expand(-1,						self.num_choices,						-1		).contiguous()
     UpperCAmelCase			:				List[str]										=					model(
         __A,						attention_mask=__A,						token_type_ids=__A,						labels=__A,						)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.num_choices)		)



   def 			__magic_name__       (				self				:	Dict,						__A				:	Tuple,						__A				:	int,						__A				:	Any,						__A				:	Optional[Any],						__A				:	Optional[Any],						__A				:	Tuple,						__A				:	str,						__A				:	List[str],						):
     UpperCAmelCase			:				Union[str, Any]										=					self.num_labels
     UpperCAmelCase			:				Optional[int]										=					FunnelForTokenClassification(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Any										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.seq_length, self.num_labels)		)



   def 			__magic_name__       (				self				:	int,						__A				:	Optional[int],						__A				:	List[str],						__A				:	List[Any],						__A				:	int,						__A				:	str,						__A				:	Any,						__A				:	Union[str, Any],						__A				:	str,						):
     UpperCAmelCase			:				str										=					FunnelForQuestionAnswering(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Optional[Any]										=					model(
         __A,						attention_mask=__A,						token_type_ids=__A,						start_positions=__A,						end_positions=__A,						)
     self.parent.assertEqual(result.start_logits.shape,						(self.batch_size, self.seq_length)		)
     self.parent.assertEqual(result.end_logits.shape,						(self.batch_size, self.seq_length)		)





   def 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase			:				Dict										=					self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		
     )			:				Optional[int]										=					config_and_inputs
     UpperCAmelCase			:				int										=					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
     return config, inputs_dict






@require_torch
class    __UpperCAmelCase	(			lowerCamelCase__     ,    lowerCamelCase__     ,    unittest.TestCase		):
   UpperCamelCase  	=	False
   UpperCamelCase  	=	False
   UpperCamelCase  	=	(
       (
           FunnelModel,
           FunnelForMaskedLM,
           FunnelForPreTraining,
           FunnelForQuestionAnswering,
           FunnelForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   UpperCamelCase  	=	(
       {
           """feature-extraction""": (FunnelBaseModel, FunnelModel),
           """fill-mask""": FunnelForMaskedLM,
           """question-answering""": FunnelForQuestionAnswering,
           """text-classification""": FunnelForSequenceClassification,
           """token-classification""": FunnelForTokenClassification,
           """zero-shot""": FunnelForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )



   def 			__magic_name__       (				self				:	Dict,						__A				:	List[Any],						__A				:	str,						__A				:	Dict=False		):
     UpperCAmelCase			:				int										=					super()._prepare_for_class(__A,						__A,						return_labels=__A		)

     if return_labels:
       if model_class in get_values(__A		):
         UpperCAmelCase			:				Tuple										=					torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length),						dtype=torch.long,						device=__A		)
     return inputs_dict



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				Optional[int]										=					FunnelModelTester(self		)
     UpperCAmelCase			:				str										=					ConfigTester(self,						config_class=__A		)



   def 			__magic_name__       (				self				:	Dict		):
     self.config_tester.run_common_tests()



   def 			__magic_name__       (				self				:	int		):
     UpperCAmelCase			:				Optional[Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A		)



   def 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase			:				List[Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_pretraining(*__A		)



   def 			__magic_name__       (				self				:	Dict		):
     UpperCAmelCase			:				str										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*__A		)



   def 			__magic_name__       (				self				:	str		):
     UpperCAmelCase			:				Any										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*__A		)



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				int										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*__A		)





   def 			__magic_name__       (				self				:	Any,						__A				:	List[Any]		):
     if hasattr(__A,						'''weight'''		) and module.weight is not None:
       module.weight.data.fill_(3		)
     if hasattr(__A,						'''bias'''		) and module.bias is not None:
       module.bias.data.fill_(3		)

     for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
       if hasattr(__A,						__A		) and getattr(__A,						__A		) is not None:
         UpperCAmelCase			:				Optional[int]										=					getattr(__A,						__A		)
         weight.data.fill_(3		)






@require_torch
class    __UpperCAmelCase	(			lowerCamelCase__     ,    unittest.TestCase		):
   UpperCamelCase  	=	False
   UpperCamelCase  	=	False
   UpperCamelCase  	=	(
       (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
   )



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				Optional[int]										=					FunnelModelTester(self,						base=__A		)
     UpperCAmelCase			:				Tuple										=					ConfigTester(self,						config_class=__A		)



   def 			__magic_name__       (				self				:	Union[str, Any]		):
     self.config_tester.run_common_tests()



   def 			__magic_name__       (				self				:	Any		):
     UpperCAmelCase			:				Union[str, Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_base_model(*__A		)



   def 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase			:				Union[str, Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*__A		)



   def 			__magic_name__       (				self				:	Optional[int]		):
     UpperCAmelCase			:				Tuple										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*__A		)



   def 			__magic_name__       (				self				:	Optional[int]		):
     UpperCAmelCase					,		UpperCAmelCase			:				List[str]										=					self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase			:				List[str]										=					True

     for model_class in self.all_model_classes:
       if model_class.__name__ == "FunnelBaseModel":
         continue
       UpperCAmelCase			:				Any										=					model_class(__A		)
       model.to(__A		)
       model.train()
       UpperCAmelCase			:				str										=					self._prepare_for_class(__A,						__A,						return_labels=__A		)
       UpperCAmelCase			:				str										=					model(**__A		).loss
       loss.backward()





   def 			__magic_name__       (				self				:	str,						__A				:	Any		):
     if hasattr(__A,						'''weight'''		) and module.weight is not None:
       module.weight.data.fill_(3		)
     if hasattr(__A,						'''bias'''		) and module.bias is not None:
       module.bias.data.fill_(3		)

     for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
       if hasattr(__A,						__A		) and getattr(__A,						__A		) is not None:
         UpperCAmelCase			:				Optional[int]										=					getattr(__A,						__A		)
         weight.data.fill_(3		)






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



   def 			__magic_name__       (				self				:	int		):
     UpperCAmelCase			:				Optional[Any]										=					1_3
     UpperCAmelCase			:				Tuple										=					7
     UpperCAmelCase			:				Optional[Any]										=					torch.arange(0,						batch_size * sequence_length		).long().reshape(__A,						__A		)
     UpperCAmelCase			:				List[Any]										=					[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
     UpperCAmelCase			:				Optional[Any]										=					torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]		)

     UpperCAmelCase			:				Dict										=					FunnelModel.from_pretrained('''sgugger/funnel-random-tiny'''		)
     UpperCAmelCase			:				List[Any]										=					model(__A,						token_type_ids=__A		)[0].abs()

     UpperCAmelCase			:				Union[str, Any]										=					torch.tensor(2_3_4_4.8_3_5_2		)
     UpperCAmelCase			:				List[Any]										=					torch.tensor(0.8_0_5_2		)
     self.assertTrue(torch.allclose(output.sum(),						__A,						atol=1E-4		)		)
     self.assertTrue(torch.allclose(output.mean(),						__A,						atol=1E-4		)		)

     UpperCAmelCase			:				Union[str, Any]										=					torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]		)
     UpperCAmelCase			:				Union[str, Any]										=					model(__A,						attention_mask=__A,						token_type_ids=__A		)[0].abs()

     UpperCAmelCase			:				str										=					torch.tensor(2_3_4_3.8_4_2_5		)
     UpperCAmelCase			:				Optional[Any]										=					torch.tensor(0.8_0_4_9		)
     self.assertTrue(torch.allclose(output.sum(),						__A,						atol=1E-4		)		)
     self.assertTrue(torch.allclose(output.mean(),						__A,						atol=1E-4		)		)





   @slow
   def 			__magic_name__       (				self				:	int		):
     UpperCAmelCase			:				List[Any]										=					FunnelTokenizer.from_pretrained('''huggingface/funnel-small'''		)
     UpperCAmelCase			:				List[str]										=					FunnelModel.from_pretrained('''huggingface/funnel-small'''		)
     UpperCAmelCase			:				List[str]										=					tokenizer('''Hello! I am the Funnel Transformer model.''',						return_tensors='''pt'''		)
     UpperCAmelCase			:				Optional[Any]										=					model(**__A		)[0]

     UpperCAmelCase			:				int										=					torch.tensor(2_3_5.7_2_4_6		)
     UpperCAmelCase			:				Optional[int]										=					torch.tensor(0.0_2_5_6		)
     self.assertTrue(torch.allclose(output.sum(),						__A,						atol=1E-4		)		)
     self.assertTrue(torch.allclose(output.mean(),						__A,						atol=1E-4		)		)





