


import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase								=	logging.get_logger(__name__)


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  A_      :				Optional[int]		     =					SwinConfig(
      embed_dim=192 ,			depths=(2, 2, 18, 2) ,			num_heads=(6, 12, 24, 48) ,			window_size=12 ,			out_features=['''stage2''', '''stage3''', '''stage4'''] ,			)

  A_      :				Optional[int]		     =					DetaConfig(
      backbone_config=SCREAMING_SNAKE_CASE ,			num_queries=900 ,			encoder_ffn_dim=2_048 ,			decoder_ffn_dim=2_048 ,			num_feature_levels=5 ,			assign_first_stage=SCREAMING_SNAKE_CASE ,			with_box_refine=SCREAMING_SNAKE_CASE ,			two_stage=SCREAMING_SNAKE_CASE ,			)

  # set labels
  A_      :				int		     =					'''huggingface/label-files'''
  if "o365" in model_name:
    A_      :				str		     =					366
    A_      :				Optional[Any]		     =					'''object365-id2label.json'''
  else:
    A_      :				Any		     =					91
    A_      :				Dict		     =					'''coco-detection-id2label.json'''

  A_      :				Tuple		     =					num_labels
  A_      :				int		     =					json.load(open(cached_download(hf_hub_url(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE ,			repo_type='''dataset'''       )       ) ,			'''r'''       )       )
  A_      :				List[Any]		     =					{int(SCREAMING_SNAKE_CASE       ): v for k, v in idalabel.items()}
  A_      :				Union[str, Any]		     =					idalabel
  A_      :				Optional[int]		     =					{v: k for k, v in idalabel.items()}

  return config


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  A_      :				List[str]		     =					[]

  # stem
  # fmt: off
  rename_keys.append(('''backbone.0.body.patch_embed.proj.weight''', '''model.backbone.model.embeddings.patch_embeddings.projection.weight''')       )
  rename_keys.append(('''backbone.0.body.patch_embed.proj.bias''', '''model.backbone.model.embeddings.patch_embeddings.projection.bias''')       )
  rename_keys.append(('''backbone.0.body.patch_embed.norm.weight''', '''model.backbone.model.embeddings.norm.weight''')       )
  rename_keys.append(('''backbone.0.body.patch_embed.norm.bias''', '''model.backbone.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.0.body.layers.{i}.blocks.{j}.norm1.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm1.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm2.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.norm2.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias''', f'''model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias''')       )

    if i < 3:
      rename_keys.append((f'''backbone.0.body.layers.{i}.downsample.reduction.weight''', f'''model.backbone.model.encoder.layers.{i}.downsample.reduction.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.downsample.norm.weight''', f'''model.backbone.model.encoder.layers.{i}.downsample.norm.weight''')       )
      rename_keys.append((f'''backbone.0.body.layers.{i}.downsample.norm.bias''', f'''model.backbone.model.encoder.layers.{i}.downsample.norm.bias''')       )

  rename_keys.append(('''backbone.0.body.norm1.weight''', '''model.backbone.model.hidden_states_norms.stage2.weight''')       )
  rename_keys.append(('''backbone.0.body.norm1.bias''', '''model.backbone.model.hidden_states_norms.stage2.bias''')       )
  rename_keys.append(('''backbone.0.body.norm2.weight''', '''model.backbone.model.hidden_states_norms.stage3.weight''')       )
  rename_keys.append(('''backbone.0.body.norm2.bias''', '''model.backbone.model.hidden_states_norms.stage3.bias''')       )
  rename_keys.append(('''backbone.0.body.norm3.weight''', '''model.backbone.model.hidden_states_norms.stage4.weight''')       )
  rename_keys.append(('''backbone.0.body.norm3.bias''', '''model.backbone.model.hidden_states_norms.stage4.bias''')       )

  # transformer encoder
  for i in range(config.encoder_layers       ):
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight''', f'''model.encoder.layers.{i}.self_attn.sampling_offsets.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias''', f'''model.encoder.layers.{i}.self_attn.sampling_offsets.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.attention_weights.weight''', f'''model.encoder.layers.{i}.self_attn.attention_weights.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.attention_weights.bias''', f'''model.encoder.layers.{i}.self_attn.attention_weights.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.value_proj.weight''', f'''model.encoder.layers.{i}.self_attn.value_proj.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.value_proj.bias''', f'''model.encoder.layers.{i}.self_attn.value_proj.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.output_proj.weight''', f'''model.encoder.layers.{i}.self_attn.output_proj.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.output_proj.bias''', f'''model.encoder.layers.{i}.self_attn.output_proj.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.norm1.weight''', f'''model.encoder.layers.{i}.self_attn_layer_norm.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.norm1.bias''', f'''model.encoder.layers.{i}.self_attn_layer_norm.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.linear1.weight''', f'''model.encoder.layers.{i}.fc1.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.linear1.bias''', f'''model.encoder.layers.{i}.fc1.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.linear2.weight''', f'''model.encoder.layers.{i}.fc2.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.linear2.bias''', f'''model.encoder.layers.{i}.fc2.bias''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.norm2.weight''', f'''model.encoder.layers.{i}.final_layer_norm.weight''')       )
    rename_keys.append((f'''transformer.encoder.layers.{i}.norm2.bias''', f'''model.encoder.layers.{i}.final_layer_norm.bias''')       )

  # transformer decoder
  for i in range(config.decoder_layers       ):
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight''', f'''model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias''', f'''model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.attention_weights.weight''', f'''model.decoder.layers.{i}.encoder_attn.attention_weights.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.attention_weights.bias''', f'''model.decoder.layers.{i}.encoder_attn.attention_weights.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.value_proj.weight''', f'''model.decoder.layers.{i}.encoder_attn.value_proj.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.value_proj.bias''', f'''model.decoder.layers.{i}.encoder_attn.value_proj.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.output_proj.weight''', f'''model.decoder.layers.{i}.encoder_attn.output_proj.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.output_proj.bias''', f'''model.decoder.layers.{i}.encoder_attn.output_proj.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.norm1.weight''', f'''model.decoder.layers.{i}.encoder_attn_layer_norm.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.norm1.bias''', f'''model.decoder.layers.{i}.encoder_attn_layer_norm.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', f'''model.decoder.layers.{i}.self_attn.out_proj.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', f'''model.decoder.layers.{i}.self_attn.out_proj.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.norm2.weight''', f'''model.decoder.layers.{i}.self_attn_layer_norm.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.norm2.bias''', f'''model.decoder.layers.{i}.self_attn_layer_norm.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.linear1.weight''', f'''model.decoder.layers.{i}.fc1.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.linear1.bias''', f'''model.decoder.layers.{i}.fc1.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.linear2.weight''', f'''model.decoder.layers.{i}.fc2.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.linear2.bias''', f'''model.decoder.layers.{i}.fc2.bias''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.norm3.weight''', f'''model.decoder.layers.{i}.final_layer_norm.weight''')       )
    rename_keys.append((f'''transformer.decoder.layers.{i}.norm3.bias''', f'''model.decoder.layers.{i}.final_layer_norm.bias''')       )

  # fmt: on

  return rename_keys


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
  A_      :				Any		     =					dct.pop(SCREAMING_SNAKE_CASE       )
  A_      :				Optional[int]		     =					val


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
  A_      :				Tuple		     =					[int(backbone_config.embed_dim * 2**i       ) for i in range(len(backbone_config.depths       )       )]
  for i in range(len(backbone_config.depths       )       ):
    A_      :				str		     =					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)
      A_      :				List[str]		     =					state_dict.pop(f'''backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight'''       )
      A_      :				List[Any]		     =					state_dict.pop(f'''backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias'''       )
      # next, add query, keys and values (in that order) to the state dict
      A_      :				Any		     =					in_proj_weight[:dim, :]
      A_      :				int		     =					in_proj_bias[: dim]
      A_      :				Union[str, Any]		     =					in_proj_weight[
          dim : dim * 2, :
      ]
      A_      :				List[str]		     =					in_proj_bias[
          dim : dim * 2
      ]
      A_      :				Tuple		     =					in_proj_weight[
          -dim :, :
      ]
      A_      :				str		     =					in_proj_bias[-dim :]
      # fmt: on


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
  # transformer decoder self-attention layers
  A_      :				Tuple		     =					config.d_model
  for i in range(config.decoder_layers       ):
    # read in weights + bias of input projection layer of self-attention
    A_      :				Optional[Any]		     =					state_dict.pop(f'''transformer.decoder.layers.{i}.self_attn.in_proj_weight'''       )
    A_      :				str		     =					state_dict.pop(f'''transformer.decoder.layers.{i}.self_attn.in_proj_bias'''       )
    # next, add query, keys and values (in that order) to the state dict
    A_      :				List[str]		     =					in_proj_weight[:hidden_size, :]
    A_      :				Optional[Any]		     =					in_proj_bias[:hidden_size]
    A_      :				Union[str, Any]		     =					in_proj_weight[
        hidden_size : hidden_size * 2, :
    ]
    A_      :				List[Any]		     =					in_proj_bias[hidden_size : hidden_size * 2]
    A_      :				Tuple		     =					in_proj_weight[-hidden_size:, :]
    A_      :				Optional[int]		     =					in_proj_bias[-hidden_size:]


def   _SCREAMING_SNAKE_CASE	( ):
  A_      :				Optional[Any]		     =					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
  A_      :				Union[str, Any]		     =					Image.open(requests.get(SCREAMING_SNAKE_CASE ,			stream=SCREAMING_SNAKE_CASE       ).raw       )

  return im


@torch.no_grad()
def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
  A_      :				Union[str, Any]		     =					get_deta_config(SCREAMING_SNAKE_CASE       )

  # load original state dict
  if model_name == "deta-swin-large":
    A_      :				List[str]		     =					hf_hub_download(repo_id='''nielsr/deta-checkpoints''' ,			filename='''adet_swin_ft.pth'''       )
  elif model_name == "deta-swin-large-o365":
    A_      :				str		     =					hf_hub_download(repo_id='''jozhang97/deta-swin-l-o365''' ,			filename='''deta_swin_pt_o365.pth'''       )
  else:
    raise ValueError(f'''Model name {model_name} not supported'''       )

  A_      :				List[str]		     =					torch.load(SCREAMING_SNAKE_CASE ,			map_location='''cpu'''       )['''model''']

  # original state dict
  for name, param in state_dict.items():
    print(SCREAMING_SNAKE_CASE ,			param.shape       )

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

  # fix some prefixes
  for key in state_dict.copy().keys():
    if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
      A_      :				List[Any]		     =					state_dict.pop(SCREAMING_SNAKE_CASE       )
      A_      :				List[Any]		     =					val
    if "input_proj" in key:
      A_      :				Dict		     =					state_dict.pop(SCREAMING_SNAKE_CASE       )
      A_      :				Dict		     =					val
    if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
      A_      :				Union[str, Any]		     =					state_dict.pop(SCREAMING_SNAKE_CASE       )
      A_      :				Optional[int]		     =					val

    # finally, create HuggingFace model and load state dict
  A_      :				Dict		     =					DetaForObjectDetection(SCREAMING_SNAKE_CASE       )
  model.load_state_dict(SCREAMING_SNAKE_CASE       )
  model.eval()

  A_      :				Optional[Any]		     =					'''cuda''' if torch.cuda.is_available() else '''cpu'''
  model.to(SCREAMING_SNAKE_CASE       )

  # load image processor
  A_      :				Optional[Any]		     =					DetaImageProcessor(format='''coco_detection'''       )

  # verify our conversion on image
  A_      :				Dict		     =					prepare_img()
  A_      :				List[str]		     =					processor(images=SCREAMING_SNAKE_CASE ,			return_tensors='''pt'''       )
  A_      :				Union[str, Any]		     =					encoding['''pixel_values''']
  A_      :				Union[str, Any]		     =					model(pixel_values.to(SCREAMING_SNAKE_CASE       )       )

  # verify logits
  print('''Logits:''' ,			outputs.logits[0, :3, :3]       )
  print('''Boxes:''' ,			outputs.pred_boxes[0, :3, :3]       )
  if model_name == "deta-swin-large":
    A_      :				Dict		     =					torch.tensor(
        [[-7.6_3_0_8, -2.8_4_8_5, -5.3_7_3_7], [-7.2_0_3_7, -4.5_5_0_5, -4.8_0_2_7], [-7.2_9_4_3, -4.2_6_1_1, -4.6_6_1_7]]       )
    A_      :				int		     =					torch.tensor([[0.4_9_8_7, 0.4_9_6_9, 0.9_9_9_9], [0.2_5_4_9, 0.5_4_9_8, 0.4_8_0_5], [0.5_4_9_8, 0.2_7_5_7, 0.0_5_6_9]]       )
  elif model_name == "deta-swin-large-o365":
    A_      :				Optional[int]		     =					torch.tensor(
        [[-8.0_1_2_2, -3.5_7_2_0, -4.9_7_1_7], [-8.1_5_4_7, -3.6_8_8_6, -4.6_3_8_9], [-7.6_6_1_0, -3.6_1_9_4, -5.0_1_3_4]]       )
    A_      :				Tuple		     =					torch.tensor([[0.2_5_2_3, 0.5_5_4_9, 0.4_8_8_1], [0.7_7_1_5, 0.4_1_4_9, 0.4_6_0_1], [0.5_5_0_3, 0.2_7_5_3, 0.0_5_7_5]]       )
  assert torch.allclose(outputs.logits[0, :3, :3] ,			expected_logits.to(SCREAMING_SNAKE_CASE       ) ,			atol=1e-4       )
  assert torch.allclose(outputs.pred_boxes[0, :3, :3] ,			expected_boxes.to(SCREAMING_SNAKE_CASE       ) ,			atol=1e-4       )
  print('''Everything ok!'''       )

  if pytorch_dump_folder_path:
    # Save model and processor
    logger.info(f'''Saving PyTorch model and processor to {pytorch_dump_folder_path}...'''       )
    Path(SCREAMING_SNAKE_CASE       ).mkdir(exist_ok=SCREAMING_SNAKE_CASE       )
    model.save_pretrained(SCREAMING_SNAKE_CASE       )
    processor.save_pretrained(SCREAMING_SNAKE_CASE       )

  # Push to hub
  if push_to_hub:
    print('''Pushing model and processor to hub...'''       )
    model.push_to_hub(f'''jozhang97/{model_name}'''       )
    processor.push_to_hub(f'''jozhang97/{model_name}'''       )


if __name__ == "__main__":
     UpperCamelCase								=	argparse.ArgumentParser()

     parser.add_argument(
         """--model_name""",
         type=str,
         default="""deta-swin-large""",
         choices=["""deta-swin-large""", """deta-swin-large-o365"""],
         help="""Name of the model you'd like to convert.""",
     )
     parser.add_argument(
         """--pytorch_dump_folder_path""",
         default=None,
         type=str,
         help="""Path to the folder to output PyTorch model.""",
     )
     parser.add_argument(
         """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
     )
     UpperCamelCase								=	parser.parse_args()
     convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.generation import (
         FlaxForcedBOSTokenLogitsProcessor,
         FlaxForcedEOSTokenLogitsProcessor,
         FlaxLogitsProcessorList,
         FlaxMinLengthLogitsProcessor,
         FlaxTemperatureLogitsWarper,
         FlaxTopKLogitsWarper,
         FlaxTopPLogitsWarper,
     )






@require_flax
class   _lowerCamelCase (					unittest.TestCase			):




       """simple docstring"""



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)->List[Any]:




         '''simple docstring'''





         A_      :				Optional[int]		     =					jnp.ones((batch_size, length)		) / length
         return scores



       def        _snake_case       (      self		)->Tuple:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					None
         A_      :				Any		     =					20

         A_      :				Any		     =					self._get_uniform_logits(batch_size=2				,       length=_SCREAMING_SNAKE_CASE		)

         # tweak scores to not be uniform anymore
         A_      :				Dict		     =					scores.at[1, 5].set((1 / length) + 0.1		)  # peak, 1st batch
         A_      :				Tuple		     =					scores.at[1, 10].set((1 / length) - 0.4		)  # valley, 1st batch

         # compute softmax
         A_      :				List[str]		     =					jax.nn.softmax(_SCREAMING_SNAKE_CASE				,       axis=-1		)

         A_      :				Any		     =					FlaxTemperatureLogitsWarper(temperature=0.5		)
         A_      :				Optional[int]		     =					FlaxTemperatureLogitsWarper(temperature=1.3		)

         A_      :				Optional[Any]		     =					jax.nn.softmax(temp_dist_warper_sharper(_SCREAMING_SNAKE_CASE				,       scores.copy()				,       cur_len=_SCREAMING_SNAKE_CASE		)				,       axis=-1		)
         A_      :				List[Any]		     =					jax.nn.softmax(temp_dist_warper_smoother(_SCREAMING_SNAKE_CASE				,       scores.copy()				,       cur_len=_SCREAMING_SNAKE_CASE		)				,       axis=-1		)

         # uniform distribution stays uniform
         self.assertTrue(jnp.allclose(probs[0, :]				,       warped_prob_sharp[0, :]				,       atol=1e-3		)		)
         self.assertTrue(jnp.allclose(probs[0, :]				,       warped_prob_smooth[0, :]				,       atol=1e-3		)		)

         # sharp peaks get higher, valleys get lower
         self.assertLess(probs[1, :].max()				,       warped_prob_sharp[1, :].max()		)
         self.assertGreater(probs[1, :].min()				,       warped_prob_sharp[1, :].min()		)

         # smooth peaks get lower, valleys get higher
         self.assertGreater(probs[1, :].max()				,       warped_prob_smooth[1, :].max()		)
         self.assertLess(probs[1, :].min()				,       warped_prob_smooth[1, :].min()		)



       def        _snake_case       (      self		)->List[Any]:




         '''simple docstring'''





         A_      :				Any		     =					None
         A_      :				List[Any]		     =					10
         A_      :				str		     =					2

         # create ramp distribution
         A_      :				Any		     =					np.broadcast_to(np.arange(_SCREAMING_SNAKE_CASE		)[None, :]				,       (batch_size, vocab_size)		).copy()
         A_      :				List[Any]		     =					ramp_logits[1:, : vocab_size // 2] + vocab_size

         A_      :				Any		     =					FlaxTopKLogitsWarper(3		)

         A_      :				Tuple		     =					top_k_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)

         # check that correct tokens are filtered
         self.assertListEqual(jnp.isinf(scores[0]		).tolist()				,       7 * [True] + 3 * [False]		)
         self.assertListEqual(jnp.isinf(scores[1]		).tolist()				,       2 * [True] + 3 * [False] + 5 * [True]		)

         # check special case
         A_      :				Optional[int]		     =					5
         A_      :				List[Any]		     =					FlaxTopKLogitsWarper(top_k=1				,       filter_value=0.0				,       min_tokens_to_keep=3		)

         A_      :				Optional[Any]		     =					np.broadcast_to(np.arange(_SCREAMING_SNAKE_CASE		)[None, :]				,       (batch_size, length)		).copy()
         A_      :				Dict		     =					top_k_warp_safety_check(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)

         # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
         self.assertListEqual((scores == 0.0).sum(axis=-1		).tolist()				,       [2, 2]		)



       def        _snake_case       (      self		)->Any:




         '''simple docstring'''





         A_      :				str		     =					None
         A_      :				Optional[Any]		     =					10
         A_      :				Any		     =					2

         # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
         A_      :				Optional[int]		     =					np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]		)		)

         A_      :				str		     =					FlaxTopPLogitsWarper(0.8		)
         A_      :				Optional[int]		     =					np.exp(top_p_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)		)

         # dist should be filtered to keep min num values so that sum is >= top_p
         # exp (-inf) => 0
         A_      :				Tuple		     =					np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]		)
         self.assertTrue(np.allclose(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       atol=1e-3		)		)

         # check edge cases with negative and extreme logits
         A_      :				Union[str, Any]		     =					np.broadcast_to(np.arange(_SCREAMING_SNAKE_CASE		)[None, :]				,       (batch_size, vocab_size)		).copy() - (
             vocab_size // 2
         )

         # make ramp_logits more extreme
         A_      :				str		     =					ramp_logits[1] * 1_0_0.0

         # make sure at least 2 tokens are kept
         A_      :				str		     =					FlaxTopPLogitsWarper(0.9				,       min_tokens_to_keep=2				,       filter_value=0.0		)
         A_      :				str		     =					top_p_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)

         # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
         self.assertListEqual((filtered_dist != 0.0).sum(axis=-1		).tolist()				,       [3, 2]		)



       def        _snake_case       (      self		)->Any:




         '''simple docstring'''





         A_      :				str		     =					20
         A_      :				Union[str, Any]		     =					4
         A_      :				Optional[Any]		     =					0

         A_      :				Union[str, Any]		     =					FlaxMinLengthLogitsProcessor(min_length=10				,       eos_token_id=_SCREAMING_SNAKE_CASE		)

         # check that min length is applied at length 5
         A_      :				int		     =					ids_tensor((batch_size, 20)				,       vocab_size=20		)
         A_      :				List[Any]		     =					5
         A_      :				Optional[int]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Optional[int]		     =					min_dist_processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()				,       4 * [-float('''inf'''		)]		)

         # check that min length is not applied anymore at length 15
         A_      :				Tuple		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Any		     =					15
         A_      :				int		     =					min_dist_processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         self.assertFalse(jnp.isinf(_SCREAMING_SNAKE_CASE		).any()		)



       def        _snake_case       (      self		)->Optional[Any]:




         '''simple docstring'''





         A_      :				Optional[int]		     =					20
         A_      :				Optional[int]		     =					4
         A_      :				Optional[int]		     =					0

         A_      :				Optional[Any]		     =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_SCREAMING_SNAKE_CASE		)

         # check that all scores are -inf except the bos_token_id score
         A_      :				Optional[Any]		     =					ids_tensor((batch_size, 1)				,       vocab_size=20		)
         A_      :				str		     =					1
         A_      :				List[str]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				List[str]		     =					logits_processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]		).all()		)
         self.assertListEqual(scores[:, bos_token_id].tolist()				,       4 * [0]		)  # score for bos_token_id shold be zero

         # check that bos_token_id is not forced if current length is greater than 1
         A_      :				Optional[int]		     =					3
         A_      :				List[Any]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Tuple		     =					logits_processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         self.assertFalse(jnp.isinf(_SCREAMING_SNAKE_CASE		).any()		)



       def        _snake_case       (      self		)->List[str]:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					20
         A_      :				str		     =					4
         A_      :				Dict		     =					0
         A_      :				Optional[int]		     =					5

         A_      :				Tuple		     =					FlaxForcedEOSTokenLogitsProcessor(max_length=_SCREAMING_SNAKE_CASE				,       eos_token_id=_SCREAMING_SNAKE_CASE		)

         # check that all scores are -inf except the eos_token_id when max_length is reached
         A_      :				List[Any]		     =					ids_tensor((batch_size, 4)				,       vocab_size=20		)
         A_      :				Any		     =					4
         A_      :				Optional[Any]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				List[Any]		     =					logits_processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]		).all()		)
         self.assertListEqual(scores[:, eos_token_id].tolist()				,       4 * [0]		)  # score for eos_token_id should be zero

         # check that eos_token_id is not forced if max_length is not reached
         A_      :				int		     =					3
         A_      :				Union[str, Any]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Dict		     =					logits_processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         self.assertFalse(jnp.isinf(_SCREAMING_SNAKE_CASE		).any()		)



       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         A_      :				str		     =					4
         A_      :				Dict		     =					10
         A_      :				Union[str, Any]		     =					15
         A_      :				str		     =					2
         A_      :				int		     =					1
         A_      :				List[str]		     =					15

         # dummy input_ids and scores
         A_      :				Tuple		     =					ids_tensor((batch_size, sequence_length)				,       _SCREAMING_SNAKE_CASE		)
         A_      :				int		     =					input_ids.copy()

         A_      :				List[Any]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Union[str, Any]		     =					scores.copy()

         # instantiate all dist processors
         A_      :				Dict		     =					FlaxTemperatureLogitsWarper(temperature=0.5		)
         A_      :				Any		     =					FlaxTopKLogitsWarper(3		)
         A_      :				List[Any]		     =					FlaxTopPLogitsWarper(0.8		)

         # instantiate all logits processors
         A_      :				List[Any]		     =					FlaxMinLengthLogitsProcessor(min_length=10				,       eos_token_id=_SCREAMING_SNAKE_CASE		)
         A_      :				Optional[int]		     =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_SCREAMING_SNAKE_CASE		)
         A_      :				Dict		     =					FlaxForcedEOSTokenLogitsProcessor(max_length=_SCREAMING_SNAKE_CASE				,       eos_token_id=_SCREAMING_SNAKE_CASE		)

         A_      :				Union[str, Any]		     =					10

         # no processor list
         A_      :				int		     =					temp_dist_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         A_      :				List[str]		     =					top_k_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         A_      :				Any		     =					top_p_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         A_      :				Dict		     =					min_dist_proc(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         A_      :				Optional[int]		     =					bos_dist_proc(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
         A_      :				Optional[Any]		     =					eos_dist_proc(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)

         # with processor list
         A_      :				Any		     =					FlaxLogitsProcessorList(
             [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
         A_      :				List[str]		     =					processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)

         # scores should be equal
         self.assertTrue(jnp.allclose(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       atol=1e-3		)		)

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()				,       input_ids_comp.tolist()		)







       def        _snake_case       (      self		)->Dict:




         '''simple docstring'''





         A_      :				str		     =					4
         A_      :				Dict		     =					10
         A_      :				Tuple		     =					15
         A_      :				List[str]		     =					2
         A_      :				List[str]		     =					1
         A_      :				Union[str, Any]		     =					15

         # dummy input_ids and scores
         A_      :				Any		     =					ids_tensor((batch_size, sequence_length)				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Union[str, Any]		     =					input_ids.copy()

         A_      :				Optional[Any]		     =					self._get_uniform_logits(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Tuple		     =					scores.copy()

         # instantiate all dist processors
         A_      :				List[str]		     =					FlaxTemperatureLogitsWarper(temperature=0.5		)
         A_      :				Optional[Any]		     =					FlaxTopKLogitsWarper(3		)
         A_      :				int		     =					FlaxTopPLogitsWarper(0.8		)

         # instantiate all logits processors
         A_      :				List[Any]		     =					FlaxMinLengthLogitsProcessor(min_length=10				,       eos_token_id=_SCREAMING_SNAKE_CASE		)
         A_      :				Dict		     =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_SCREAMING_SNAKE_CASE		)
         A_      :				Optional[Any]		     =					FlaxForcedEOSTokenLogitsProcessor(max_length=_SCREAMING_SNAKE_CASE				,       eos_token_id=_SCREAMING_SNAKE_CASE		)

         A_      :				str		     =					10

         # no processor list
         def run_no_processor_list(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		):
           A_      :				int		     =					temp_dist_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           A_      :				Optional[Any]		     =					top_k_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           A_      :				List[Any]		     =					top_p_warp(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           A_      :				Dict		     =					min_dist_proc(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           A_      :				Any		     =					bos_dist_proc(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           A_      :				Optional[Any]		     =					eos_dist_proc(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           return scores

         # with processor list
         def run_processor_list(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		):
           A_      :				Optional[int]		     =					FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
           A_      :				Optional[int]		     =					processor(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       cur_len=_SCREAMING_SNAKE_CASE		)
           return scores

         A_      :				Optional[int]		     =					jax.jit(_SCREAMING_SNAKE_CASE		)
         A_      :				Union[str, Any]		     =					jax.jit(_SCREAMING_SNAKE_CASE		)

         A_      :				Dict		     =					jitted_run_no_processor_list(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				List[Any]		     =					jitted_run_processor_list(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)

         # scores should be equal
         self.assertTrue(jnp.allclose(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       atol=1e-3		)		)

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()				,       input_ids_comp.tolist()		)




