'''simple docstring'''
import argparse
import json
import pickle
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
lowercase__  :      str        =					logging.get_logger(__name__)



def  __lowerCamelCase       (      _UpperCamelCase  :  str					):


      '''simple docstring'''


      UpperCAmelCase_            =  SwinConfig.from_pretrained(
          '''microsoft/swin-tiny-patch4-window7-224'''							,    out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']					)
      UpperCAmelCase_            =  MaskFormerConfig(backbone_config=_UpperCamelCase					)

      UpperCAmelCase_            =  '''huggingface/label-files'''
      if "ade20k-full" in model_name:
            # this should be ok
            UpperCAmelCase_            =  847
            UpperCAmelCase_            =  '''maskformer-ade20k-full-id2label.json'''
      elif "ade" in model_name:
            # this should be ok
            UpperCAmelCase_            =  150
            UpperCAmelCase_            =  '''ade20k-id2label.json'''
      elif "coco-stuff" in model_name:
            # this should be ok
            UpperCAmelCase_            =  171
            UpperCAmelCase_            =  '''maskformer-coco-stuff-id2label.json'''
      elif "coco" in model_name:
            # TODO
            UpperCAmelCase_            =  133
            UpperCAmelCase_            =  '''coco-panoptic-id2label.json'''
      elif "cityscapes" in model_name:
            # this should be ok
            UpperCAmelCase_            =  19
            UpperCAmelCase_            =  '''cityscapes-id2label.json'''
      elif "vistas" in model_name:
            # this should be ok
            UpperCAmelCase_            =  65
            UpperCAmelCase_            =  '''mapillary-vistas-id2label.json'''

      UpperCAmelCase_            =  json.load(open(hf_hub_download(_UpperCamelCase							,    _UpperCamelCase							,    repo_type='''dataset'''					)							,    '''r'''					)					)
      UpperCAmelCase_            =  {int(_UpperCamelCase					): v for k, v in idalabel.items()}

      return config



def  __lowerCamelCase       (      _UpperCamelCase  :  Any					):


      '''simple docstring'''


      UpperCAmelCase_            =  []
      # stem
      # fmt: off
      rename_keys.append(('''backbone.patch_embed.proj.weight''', '''model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.weight''')					)
      rename_keys.append(('''backbone.patch_embed.proj.bias''', '''model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.bias''')					)
      rename_keys.append(('''backbone.patch_embed.norm.weight''', '''model.pixel_level_module.encoder.model.embeddings.norm.weight''')					)
      rename_keys.append(('''backbone.patch_embed.norm.bias''', '''model.pixel_level_module.encoder.model.embeddings.norm.bias''')					)
      # stages
      for i in range(len(config.backbone_config.depths					)					):
            for j in range(config.backbone_config.depths[i]					):
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.norm1.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.norm1.bias""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.attn.relative_position_bias_table""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.attn.relative_position_index""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.attn.proj.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.attn.proj.bias""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.norm2.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.norm2.bias""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.mlp.fc1.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.mlp.fc1.bias""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.mlp.fc2.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.blocks.{j}.mlp.fc2.bias""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.bias""")					)

            if i < 3:
                  rename_keys.append((F"""backbone.layers.{i}.downsample.reduction.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.reduction.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.downsample.norm.weight""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.weight""")					)
                  rename_keys.append((F"""backbone.layers.{i}.downsample.norm.bias""", F"""model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.bias""")					)
            rename_keys.append((F"""backbone.norm{i}.weight""", F"""model.pixel_level_module.encoder.hidden_states_norms.{i}.weight""")					)
            rename_keys.append((F"""backbone.norm{i}.bias""", F"""model.pixel_level_module.encoder.hidden_states_norms.{i}.bias""")					)

      # FPN
      rename_keys.append(('''sem_seg_head.layer_4.weight''', '''model.pixel_level_module.decoder.fpn.stem.0.weight''')					)
      rename_keys.append(('''sem_seg_head.layer_4.norm.weight''', '''model.pixel_level_module.decoder.fpn.stem.1.weight''')					)
      rename_keys.append(('''sem_seg_head.layer_4.norm.bias''', '''model.pixel_level_module.decoder.fpn.stem.1.bias''')					)
      for source_index, target_index in zip(range(3							,    0							,    -1					)							,    range(0							,    3					)					):
            rename_keys.append((F"""sem_seg_head.adapter_{source_index}.weight""", F"""model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight""")					)
            rename_keys.append((F"""sem_seg_head.adapter_{source_index}.norm.weight""", F"""model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight""")					)
            rename_keys.append((F"""sem_seg_head.adapter_{source_index}.norm.bias""", F"""model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias""")					)
            rename_keys.append((F"""sem_seg_head.layer_{source_index}.weight""", F"""model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight""")					)
            rename_keys.append((F"""sem_seg_head.layer_{source_index}.norm.weight""", F"""model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight""")					)
            rename_keys.append((F"""sem_seg_head.layer_{source_index}.norm.bias""", F"""model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias""")					)
      rename_keys.append(('''sem_seg_head.mask_features.weight''', '''model.pixel_level_module.decoder.mask_projection.weight''')					)
      rename_keys.append(('''sem_seg_head.mask_features.bias''', '''model.pixel_level_module.decoder.mask_projection.bias''')					)

      # Transformer decoder
      for idx in range(config.decoder_config.decoder_layers					):
            # self-attention out projection
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight""", F"""model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias""", F"""model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias""")					)
            # cross-attention out projection
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight""", F"""model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias""", F"""model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias""")					)
            # MLP 1
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight""", F"""model.transformer_module.decoder.layers.{idx}.fc1.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias""", F"""model.transformer_module.decoder.layers.{idx}.fc1.bias""")					)
            # MLP 2
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight""", F"""model.transformer_module.decoder.layers.{idx}.fc2.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias""", F"""model.transformer_module.decoder.layers.{idx}.fc2.bias""")					)
            # layernorm 1 (self-attention layernorm)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight""", F"""model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias""", F"""model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias""")					)
            # layernorm 2 (cross-attention layernorm)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight""", F"""model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias""", F"""model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias""")					)
            # layernorm 3 (final layernorm)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight""", F"""model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias""", F"""model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias""")					)

      rename_keys.append(('''sem_seg_head.predictor.transformer.decoder.norm.weight''', '''model.transformer_module.decoder.layernorm.weight''')					)
      rename_keys.append(('''sem_seg_head.predictor.transformer.decoder.norm.bias''', '''model.transformer_module.decoder.layernorm.bias''')					)

      # heads on top
      rename_keys.append(('''sem_seg_head.predictor.query_embed.weight''', '''model.transformer_module.queries_embedder.weight''')					)

      rename_keys.append(('''sem_seg_head.predictor.input_proj.weight''', '''model.transformer_module.input_projection.weight''')					)
      rename_keys.append(('''sem_seg_head.predictor.input_proj.bias''', '''model.transformer_module.input_projection.bias''')					)

      rename_keys.append(('''sem_seg_head.predictor.class_embed.weight''', '''class_predictor.weight''')					)
      rename_keys.append(('''sem_seg_head.predictor.class_embed.bias''', '''class_predictor.bias''')					)

      for i in range(3					):
            rename_keys.append((F"""sem_seg_head.predictor.mask_embed.layers.{i}.weight""", F"""mask_embedder.{i}.0.weight""")					)
            rename_keys.append((F"""sem_seg_head.predictor.mask_embed.layers.{i}.bias""", F"""mask_embedder.{i}.0.bias""")					)
      # fmt: on

      return rename_keys



def  __lowerCamelCase       (      _UpperCamelCase  :  Any							,    _UpperCamelCase  :  Union[str, Any]							,    _UpperCamelCase  :  List[str]					):


      '''simple docstring'''


      UpperCAmelCase_            =  dct.pop(_UpperCamelCase					)
      UpperCAmelCase_            =  val



def  __lowerCamelCase       (      _UpperCamelCase  :  Any							,    _UpperCamelCase  :  int					):


      '''simple docstring'''


      UpperCAmelCase_            =  [int(backbone_config.embed_dim * 2**i					) for i in range(len(backbone_config.depths					)					)]
      for i in range(len(backbone_config.depths					)					):
            UpperCAmelCase_            =  num_features[i]
            for j in range(backbone_config.depths[i]					):
                  # fmt: off
                  # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
                  UpperCAmelCase_            =  state_dict.pop(F"""backbone.layers.{i}.blocks.{j}.attn.qkv.weight"""					)
                  UpperCAmelCase_            =  state_dict.pop(F"""backbone.layers.{i}.blocks.{j}.attn.qkv.bias"""					)
                  # next, add query, keys and values (in that order) to the state dict
                  UpperCAmelCase_            =  in_proj_weight[:dim, :]
                  UpperCAmelCase_            =  in_proj_bias[: dim]
                  UpperCAmelCase_            =  in_proj_weight[
                      dim : dim * 2, :
                  ]
                  UpperCAmelCase_            =  in_proj_bias[
                      dim : dim * 2
                  ]
                  UpperCAmelCase_            =  in_proj_weight[
                      -dim :, :
                  ]
                  UpperCAmelCase_            =  in_proj_bias[-dim :]
                  # fmt: on



def  __lowerCamelCase       (      _UpperCamelCase  :  str							,    _UpperCamelCase  :  int					):


      '''simple docstring'''


      UpperCAmelCase_            =  config.decoder_config.hidden_size
      for idx in range(config.decoder_config.decoder_layers					):
            # read in weights + bias of self-attention input projection layer (in the original implementation, this is a single matrix + bias)
            UpperCAmelCase_            =  state_dict.pop(F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight"""					)
            UpperCAmelCase_            =  state_dict.pop(F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_bias"""					)
            # next, add query, keys and values (in that order) to the state dict
            UpperCAmelCase_            =  in_proj_weight[: hidden_size, :]
            UpperCAmelCase_            =  in_proj_bias[:config.hidden_size]
            UpperCAmelCase_            =  in_proj_weight[hidden_size : hidden_size * 2, :]
            UpperCAmelCase_            =  in_proj_bias[hidden_size : hidden_size * 2]
            UpperCAmelCase_            =  in_proj_weight[-hidden_size :, :]
            UpperCAmelCase_            =  in_proj_bias[-hidden_size :]
            # read in weights + bias of cross-attention input projection layer (in the original implementation, this is a single matrix + bias)
            UpperCAmelCase_            =  state_dict.pop(F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight"""					)
            UpperCAmelCase_            =  state_dict.pop(F"""sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_bias"""					)
            # next, add query, keys and values (in that order) to the state dict
            UpperCAmelCase_            =  in_proj_weight[: hidden_size, :]
            UpperCAmelCase_            =  in_proj_bias[:config.hidden_size]
            UpperCAmelCase_            =  in_proj_weight[hidden_size : hidden_size * 2, :]
            UpperCAmelCase_            =  in_proj_bias[hidden_size : hidden_size * 2]
            UpperCAmelCase_            =  in_proj_weight[-hidden_size :, :]
            UpperCAmelCase_            =  in_proj_bias[-hidden_size :]
      # fmt: on



def  __lowerCamelCase       (      ):


      '''simple docstring'''


      UpperCAmelCase_            =  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
      UpperCAmelCase_            =  Image.open(requests.get(_UpperCamelCase							,    stream=_UpperCamelCase					).raw					)
      return im



@torch.no_grad()
def  __lowerCamelCase       (      _UpperCamelCase  :  str							,    _UpperCamelCase  :  str							,    _UpperCamelCase  :  str							,    _UpperCamelCase  :  bool = False					):


      '''simple docstring'''


      UpperCAmelCase_            =  get_maskformer_config(_UpperCamelCase					)

      # load original state_dict
      with open(_UpperCamelCase							,    '''rb'''					) as f:
            UpperCAmelCase_            =  pickle.load(_UpperCamelCase					)
      UpperCAmelCase_            =  data['''model''']

      # for name, param in state_dict.items():
      #     print(name, param.shape)

      # rename keys
      UpperCAmelCase_            =  create_rename_keys(_UpperCamelCase					)
      for src, dest in rename_keys:
            rename_key(_UpperCamelCase							,    _UpperCamelCase							,    _UpperCamelCase					)
      read_in_swin_q_k_v(_UpperCamelCase							,    config.backbone_config					)
      read_in_decoder_q_k_v(_UpperCamelCase							,    _UpperCamelCase					)

      # update to torch tensors
      for key, value in state_dict.items():
            UpperCAmelCase_            =  torch.from_numpy(_UpperCamelCase					)

      # load 🤗 model
      UpperCAmelCase_            =  MaskFormerForInstanceSegmentation(_UpperCamelCase					)
      model.eval()

      for name, param in model.named_parameters():
            print(_UpperCamelCase							,    param.shape					)

      UpperCAmelCase_      ,		UpperCAmelCase_            =  model.load_state_dict(_UpperCamelCase							,    strict=_UpperCamelCase					)
      assert missing_keys == [
          "model.pixel_level_module.encoder.model.layernorm.weight",
          "model.pixel_level_module.encoder.model.layernorm.bias",
      ]
      assert len(_UpperCamelCase					) == 0, F"""Unexpected keys: {unexpected_keys}"""

      # verify results
      UpperCAmelCase_            =  prepare_img()
      if "vistas" in model_name:
            UpperCAmelCase_            =  65
      elif "cityscapes" in model_name:
            UpperCAmelCase_            =  6_5535
      else:
            UpperCAmelCase_            =  255
      UpperCAmelCase_            =  True if '''ade''' in model_name else False
      UpperCAmelCase_            =  MaskFormerImageProcessor(ignore_index=_UpperCamelCase							,    reduce_labels=_UpperCamelCase					)

      UpperCAmelCase_            =  image_processor(_UpperCamelCase							,    return_tensors='''pt'''					)

      UpperCAmelCase_            =  model(**_UpperCamelCase					)

      print('''Logits:'''							,    outputs.class_queries_logits[0, :3, :3]					)

      if model_name == "maskformer-swin-tiny-ade":
            UpperCAmelCase_            =  torch.tensor(
                [[3.6_353, -4.4_770, -2.6_065], [0.5_081, -4.2_394, -3.5_343], [2.1_909, -5.0_353, -1.9_323]]					)
      assert torch.allclose(outputs.class_queries_logits[0, :3, :3]							,    _UpperCamelCase							,    atol=1E-4					)
      print('''Looks ok!'''					)

      if pytorch_dump_folder_path is not None:
            print(F"""Saving model and image processor to {pytorch_dump_folder_path}"""					)
            Path(_UpperCamelCase					).mkdir(exist_ok=_UpperCamelCase					)
            model.save_pretrained(_UpperCamelCase					)
            image_processor.save_pretrained(_UpperCamelCase					)

      if push_to_hub:
            print('''Pushing model and image processor to the hub...'''					)
            model.push_to_hub(F"""nielsr/{model_name}"""					)
            image_processor.push_to_hub(F"""nielsr/{model_name}"""					)


if __name__ == "__main__":
      lowercase__  :      Dict        =					argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--model_name",
          default="maskformer-swin-tiny-ade",
          type=str,
          help=("Name of the MaskFormer model you'd like to convert",),
      )
      parser.add_argument(
          "--checkpoint_path",
          default="/Users/nielsrogge/Documents/MaskFormer_checkpoints/MaskFormer-Swin-tiny-ADE20k/model.pkl",
          type=str,
          help="Path to the original state dict (.pth file).",
      )
      parser.add_argument(
          "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
      )
      parser.add_argument(
          "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
      )

      lowercase__  :      Any        =					parser.parse_args()
      convert_maskformer_checkpoint(
          args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
      )






'''simple docstring'''
from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          AutoTokenizer,
          TFAutoModelForCausalLM,
          TFAutoModelForSeqaSeqLM,
          TFAutoModelForSpeechSeqaSeq,
          TFAutoModelForVisionaSeq,
          TFBartForConditionalGeneration,
          TFLogitsProcessorList,
          TFMinLengthLogitsProcessor,
          tf_top_k_top_p_filtering,
      )

if is_tensorflow_text_available():
      import tensorflow_text as text







@require_tf
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Dict  )				->List[str]:
        UpperCAmelCase_            =  tf.convert_to_tensor(
            [
                [
                    8.222_0991,  # 3rd highest value; idx. 0
                    -0.562_0044,
                    5.2322_9752,
                    4.038_6393,
                    -6.879_8378,
                    -0.5478_5802,
                    -3.201_2153,
                    2.9277_7176,
                    1.8817_1953,
                    7.3534_1276,  # 5th highest value; idx. 9
                    8.4320_7833,  # 2nd highest value; idx. 10
                    -9.8571_1836,
                    -5.9620_9236,
                    -1.1303_9161,
                    -7.111_5294,
                    -0.836_9633,
                    -5.318_6408,
                    7.0642_7407,
                    0.8136_9344,
                    -0.8202_3817,
                    -5.917_9796,
                    0.5881_3443,
                    -6.9977_8438,
                    4.7155_1189,
                    -0.1877_1637,
                    7.4402_0759,  # 4th highest value; idx. 25
                    9.3845_0987,  # 1st highest value; idx. 26
                    2.1266_2941,
                    -9.3256_2038,
                    2.3565_2522,
                ],  # cummulative prob of 5 highest values <= 0.6
                [
                    0.5842_5518,
                    4.5313_9238,
                    -5.5751_0464,
                    -6.2803_0699,
                    -7.1952_9503,
                    -4.0212_2551,
                    1.3933_7037,
                    -6.0670_7057,
                    1.5948_0517,
                    -9.64_3119,
                    0.0390_7799,
                    0.6723_1762,
                    -8.8820_6726,
                    6.2711_5922,  # 4th highest value; idx. 13
                    2.2852_0723,
                    4.8276_7506,
                    4.3042_1368,
                    8.827_5313,  # 2nd highest value; idx. 17
                    5.4402_9958,  # 5th highest value; idx. 18
                    -4.473_5794,
                    7.3857_9536,  # 3rd highest value; idx. 20
                    -2.9105_1663,
                    2.6194_6077,
                    -2.567_4762,
                    -9.4895_9302,
                    -4.0292_2645,
                    -1.3541_6918,
                    9.6770_2323,  # 1st highest value; idx. 27
                    -5.8947_8553,
                    1.8537_0467,
                ],  # cummulative prob of 5 highest values <= 0.6
            ]  ,							dtype=tf.floataa  ,							)

        UpperCAmelCase_            =  tf.convert_to_tensor(
            [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]  ,							dtype=tf.intaa  ,							)  # expected non filtered idx as noted above

        UpperCAmelCase_            =  tf.convert_to_tensor(
            [8.22_2099, 7.353_4126, 8.43_2078, 7.440_2075, 9.3_8451, 6.27_1159, 8.82_7531, 5.440_2995, 7.385_7956, 9.67_7023]  ,							dtype=tf.floataa  ,							)  # expected non filtered values as noted above

        UpperCAmelCase_            =  tf_top_k_top_p_filtering(UpperCAmelCase__  ,							top_k=10  ,							top_p=0.6  ,							min_tokens_to_keep=4  )

        UpperCAmelCase_            =  output[output != -float('''inf'''  )]
        UpperCAmelCase_            =  tf.cast(
            tf.where(tf.not_equal(UpperCAmelCase__  ,							tf.constant(-float('''inf'''  )  ,							dtype=tf.floataa  )  )  )  ,							dtype=tf.intaa  ,							)

        tf.debugging.assert_near(UpperCAmelCase__  ,							UpperCAmelCase__  ,							rtol=1e-12  )
        tf.debugging.assert_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )







@require_tf
class lowerCamelCase							(	unittest.TestCase     ,		lowerCamelCase					):




  '''simple docstring'''



  if is_tf_available():
    lowerCAmelCase__										=	{
        '''AutoModelForCausalLM''': TFAutoModelForCausalLM,
        '''AutoModelForSpeechSeq2Seq''': TFAutoModelForSpeechSeqaSeq,
        '''AutoModelForSeq2SeqLM''': TFAutoModelForSeqaSeqLM,
        '''AutoModelForVision2Seq''': TFAutoModelForVisionaSeq,
        '''LogitsProcessorList''': TFLogitsProcessorList,
        '''MinLengthLogitsProcessor''': TFMinLengthLogitsProcessor,
        '''create_tensor_fn''': tf.convert_to_tensor,
        '''floats_tensor''': floats_tensor,
        '''return_tensors''': '''tf''',
    }






  @slow
  def  lowerCAmelCase__ (      self			:       Dict  )				->Optional[Any]:
        # TF-only test: tf.saved_model export
        UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  2







        class lowerCamelCase							(	tf.Module					):




          '''simple docstring'''



          def __init__(      self			:       List[str]  ,							UpperCAmelCase__			:       List[str]  )				->Dict:
                super(UpperCAmelCase__  ,							self  ).__init__()
                UpperCAmelCase_            =  model






          @tf.function(
              input_signature=(
                  tf.TensorSpec((None, input_length)  ,							tf.intaa  ,							name='''input_ids'''  ),
                  tf.TensorSpec((None, input_length)  ,							tf.intaa  ,							name='''attention_mask'''  ),
              )  ,							jit_compile=UpperCAmelCase__  ,							)
          def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       Union[str, Any]  ,							UpperCAmelCase__			:       List[str]  )				->int:
                UpperCAmelCase_            =  self.model.generate(
                    input_ids=UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  ,							return_dict_in_generate=UpperCAmelCase__  ,							)
                return {"sequences": outputs["sequences"]}






        UpperCAmelCase_            =  [[2, 0], [102, 103]]
        UpperCAmelCase_            =  [[1, 0], [1, 1]]
        UpperCAmelCase_            =  DummyModel(model=UpperCAmelCase__  )
        with tempfile.TemporaryDirectory() as tmp_dir:
              tf.saved_model.save(UpperCAmelCase__  ,							UpperCAmelCase__  ,							signatures={'''serving_default''': dummy_model.serving}  )
              UpperCAmelCase_            =  tf.saved_model.load(UpperCAmelCase__  ).signatures['''serving_default''']
              for batch_size in range(1  ,							len(UpperCAmelCase__  ) + 1  ):
                    UpperCAmelCase_            =  {
                        '''input_ids''': tf.constant(dummy_input_ids[:batch_size]  ),
                        '''attention_mask''': tf.constant(dummy_attention_masks[:batch_size]  ),
                    }
                    UpperCAmelCase_            =  serving_func(**UpperCAmelCase__  )['''sequences''']
                    UpperCAmelCase_            =  test_model.generate(**UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  )
                    tf.debugging.assert_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       Tuple  )				->Optional[int]:
        # TF-only test: tf.saved_model export
        UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )
        UpperCAmelCase_            =  1
        UpperCAmelCase_            =  2







        class lowerCamelCase							(	tf.Module					):




          '''simple docstring'''



          def __init__(      self			:       int  ,							UpperCAmelCase__			:       int  )				->List[str]:
                super(UpperCAmelCase__  ,							self  ).__init__()
                UpperCAmelCase_            =  model






          @tf.function(
              input_signature=(
                  tf.TensorSpec((batch_size, None)  ,							tf.intaa  ,							name='''input_ids'''  ),
                  tf.TensorSpec((batch_size, None)  ,							tf.intaa  ,							name='''attention_mask'''  ),
              )  ,							jit_compile=UpperCAmelCase__  ,							)
          def  lowerCAmelCase__ (      self			:       Dict  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       Optional[Any]  )				->int:
                UpperCAmelCase_            =  self.model.generate(
                    input_ids=UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  ,							return_dict_in_generate=UpperCAmelCase__  ,							)
                return {"sequences": outputs["sequences"]}






        UpperCAmelCase_            =  [[2], [102, 103]]
        UpperCAmelCase_            =  [[1], [1, 1]]
        UpperCAmelCase_            =  DummyModel(model=UpperCAmelCase__  )
        with tempfile.TemporaryDirectory() as tmp_dir:
              tf.saved_model.save(UpperCAmelCase__  ,							UpperCAmelCase__  ,							signatures={'''serving_default''': dummy_model.serving}  )
              UpperCAmelCase_            =  tf.saved_model.load(UpperCAmelCase__  ).signatures['''serving_default''']
              for input_row in range(len(UpperCAmelCase__  )  ):
                    UpperCAmelCase_            =  {
                        '''input_ids''': tf.constant([dummy_input_ids[input_row]]  ),
                        '''attention_mask''': tf.constant([dummy_attention_masks[input_row]]  ),
                    }
                    UpperCAmelCase_            =  serving_func(**UpperCAmelCase__  )['''sequences''']
                    UpperCAmelCase_            =  test_model.generate(**UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  )
                    tf.debugging.assert_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  @require_tensorflow_text
  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->int:
        # TF-only test: tf.saved_model export
        with tempfile.TemporaryDirectory() as tmp_dir:
              # file needed to load the TF tokenizer
              hf_hub_download(repo_id='''google/flan-t5-small'''  ,							filename='''spiece.model'''  ,							local_dir=UpperCAmelCase__  )







              class lowerCamelCase							(	tf.keras.layers.Layer					):




                '''simple docstring'''



                def __init__(      self			:       List[str]  )				->Any:
                      super().__init__()
                      UpperCAmelCase_            =  text.SentencepieceTokenizer(
                          model=tf.io.gfile.GFile(os.path.join(UpperCAmelCase__  ,							'''spiece.model'''  )  ,							'''rb'''  ).read()  )
                      UpperCAmelCase_            =  TFAutoModelForSeqaSeqLM.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )






                def  lowerCAmelCase__ (      self			:       Any  ,							UpperCAmelCase__			:       Union[str, Any]  ,							*UpperCAmelCase__			:       Union[str, Any]  ,							**UpperCAmelCase__			:       str  )				->List[str]:
                      UpperCAmelCase_            =  self.tokenizer.tokenize(UpperCAmelCase__  )
                      UpperCAmelCase_      ,		UpperCAmelCase_            =  text.pad_model_inputs(
                          UpperCAmelCase__  ,							max_seq_length=64  ,							pad_value=self.model.config.pad_token_id  )
                      UpperCAmelCase_            =  self.model.generate(input_ids=UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  )
                      return self.tokenizer.detokenize(UpperCAmelCase__  )






              UpperCAmelCase_            =  CompleteSentenceTransformer()
              UpperCAmelCase_            =  tf.keras.layers.Input(shape=(1,)  ,							dtype=tf.string  ,							name='''inputs'''  )
              UpperCAmelCase_            =  complete_model(UpperCAmelCase__  )
              UpperCAmelCase_            =  tf.keras.Model(UpperCAmelCase__  ,							UpperCAmelCase__  )
              keras_model.save(UpperCAmelCase__  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->Tuple:
        # Has PT equivalent: this test relies on random sampling
        UpperCAmelCase_            =  {
            '''do_sample''': True,
            '''num_beams''': 1,
            '''top_p''': 0.7,
            '''top_k''': 10,
            '''temperature''': 0.7,
        }
        UpperCAmelCase_            =  14

        UpperCAmelCase_            =  AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )
        UpperCAmelCase_            =  '''Hello, my dog is cute and'''
        UpperCAmelCase_            =  tokenizer(UpperCAmelCase__  ,							return_tensors='''tf'''  )
        UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )

        UpperCAmelCase_            =  638
        # forces the generation to happen on CPU, to avoid GPU-related quirks
        with tf.device(''':/CPU:0'''  ):
              tf.random.set_seed(0  )
              UpperCAmelCase_            =  model.generate(**UpperCAmelCase__  ,							eos_token_id=UpperCAmelCase__  ,							**UpperCAmelCase__  )
        self.assertTrue(expectation == len(generated_tokens[0]  )  )

        UpperCAmelCase_            =  [638, 198]
        with tf.device(''':/CPU:0'''  ):
              tf.random.set_seed(0  )
              UpperCAmelCase_            =  model.generate(**UpperCAmelCase__  ,							eos_token_id=UpperCAmelCase__  ,							**UpperCAmelCase__  )
        self.assertTrue(expectation == len(generated_tokens[0]  )  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->List[str]:
        # Has PT equivalent: ample use of framework-specific code
        UpperCAmelCase_            =  AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bart'''  )
        UpperCAmelCase_            =  '''Hugging Face is a technology company based in New York and Paris.'''
        UpperCAmelCase_            =  bart_tokenizer(UpperCAmelCase__  ,							return_tensors='''tf'''  ).input_ids
        UpperCAmelCase_            =  TFBartForConditionalGeneration.from_pretrained('''hf-internal-testing/tiny-random-bart'''  )
        UpperCAmelCase_            =  bart_model.generate(UpperCAmelCase__  ).numpy()







        class lowerCamelCase							(	lowerCamelCase					):




          '''simple docstring'''



          def  lowerCAmelCase__ (      self			:       int  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       List[Any]=None  ,							**UpperCAmelCase__			:       int  )				->List[str]:
                return super().call(UpperCAmelCase__  ,							**UpperCAmelCase__  )






        UpperCAmelCase_            =  FakeBart.from_pretrained('''hf-internal-testing/tiny-random-bart'''  )
        UpperCAmelCase_            =  bart_model.generate(UpperCAmelCase__  ,							foo='''bar'''  ).numpy()
        self.assertTrue(np.array_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )  )







        class lowerCamelCase							(	bart_model.model.encoder.__class__					):




          '''simple docstring'''



          def  lowerCAmelCase__ (      self			:       str  ,							UpperCAmelCase__			:       List[str]  ,							**UpperCAmelCase__			:       Dict  )				->Any:
                return super().call(UpperCAmelCase__  ,							**UpperCAmelCase__  )






        UpperCAmelCase_            =  FakeEncoder(bart_model.config  ,							bart_model.model.shared  )
        UpperCAmelCase_            =  fake_encoder

        # Normal generation still works (the output will be different because the encoder weights are different)
        UpperCAmelCase_            =  bart_model.generate(UpperCAmelCase__  ).numpy()
        with self.assertRaises(UpperCAmelCase__  ):
              # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
              bart_model.generate(UpperCAmelCase__  ,							foo='''bar'''  )





