


'''simple docstring'''



import argparse
import json

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

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[Any]							)     ->		str:
     __lowerCamelCase   :	Dict								=		3_84
     if "tiny" in model_name:
          __lowerCamelCase   :	Optional[int]								=		[3, 3, 9, 3]
          __lowerCamelCase   :	Tuple								=		[96, 1_92, 3_84, 7_68]
     if "small" in model_name:
          __lowerCamelCase   :	Optional[Any]								=		[3, 3, 27, 3]
          __lowerCamelCase   :	int								=		[96, 1_92, 3_84, 7_68]
     if "base" in model_name:
          __lowerCamelCase   :	Union[str, Any]								=		[3, 3, 27, 3]
          __lowerCamelCase   :	str								=		[1_28, 2_56, 5_12, 10_24]
          __lowerCamelCase   :	Optional[int]								=		5_12
     if "large" in model_name:
          __lowerCamelCase   :	Union[str, Any]								=		[3, 3, 27, 3]
          __lowerCamelCase   :	List[str]								=		[1_92, 3_84, 7_68, 15_36]
          __lowerCamelCase   :	Union[str, Any]								=		7_68
     if "xlarge" in model_name:
          __lowerCamelCase   :	Optional[Any]								=		[3, 3, 27, 3]
          __lowerCamelCase   :	str								=		[2_56, 5_12, 10_24, 20_48]
          __lowerCamelCase   :	List[Any]								=		10_24

     # set label information
     __lowerCamelCase   :	Union[str, Any]								=		1_50
     __lowerCamelCase   :	List[str]								=		'huggingface/label-files'
     __lowerCamelCase   :	Optional[int]								=		'ade20k-id2label.json'
     __lowerCamelCase   :	List[Any]								=		json.load(open(hf_hub_download(UpperCAmelCase_      ,    UpperCAmelCase_      ,    repo_type='dataset'							)      ,    'r'							)							)
     __lowerCamelCase   :	Any								=		{int(UpperCAmelCase_							): v for k, v in idalabel.items()}
     __lowerCamelCase   :	Union[str, Any]								=		{v: k for k, v in idalabel.items()}

     __lowerCamelCase   :	Any								=		ConvNextConfig(
         depths=UpperCAmelCase_      ,    hidden_sizes=UpperCAmelCase_      ,    out_features=['stage1', 'stage2', 'stage3', 'stage4']							)
     __lowerCamelCase   :	Any								=		UperNetConfig(
         backbone_config=UpperCAmelCase_      ,    auxiliary_in_channels=UpperCAmelCase_      ,    num_labels=UpperCAmelCase_      ,    idalabel=UpperCAmelCase_      ,    labelaid=UpperCAmelCase_      ,    )

     return config



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Tuple							)     ->		str:
     __lowerCamelCase   :	Dict								=		[]

     # fmt: off
     # stem
     rename_keys.append(('backbone.downsample_layers.0.0.weight', 'backbone.embeddings.patch_embeddings.weight')							)
     rename_keys.append(('backbone.downsample_layers.0.0.bias', 'backbone.embeddings.patch_embeddings.bias')							)
     rename_keys.append(('backbone.downsample_layers.0.1.weight', 'backbone.embeddings.layernorm.weight')							)
     rename_keys.append(('backbone.downsample_layers.0.1.bias', 'backbone.embeddings.layernorm.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.stages.{i}.{j}.gamma', F'backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.depthwise_conv.weight', F'backbone.encoder.stages.{i}.layers.{j}.dwconv.weight')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.depthwise_conv.bias', F'backbone.encoder.stages.{i}.layers.{j}.dwconv.bias')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.norm.weight', F'backbone.encoder.stages.{i}.layers.{j}.layernorm.weight')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.norm.bias', F'backbone.encoder.stages.{i}.layers.{j}.layernorm.bias')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv1.weight', F'backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv1.bias', F'backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv2.weight', F'backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight')							)
               rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv2.bias', F'backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias')							)
          if i > 0:
               rename_keys.append((F'backbone.downsample_layers.{i}.0.weight', F'backbone.encoder.stages.{i}.downsampling_layer.0.weight')							)
               rename_keys.append((F'backbone.downsample_layers.{i}.0.bias', F'backbone.encoder.stages.{i}.downsampling_layer.0.bias')							)
               rename_keys.append((F'backbone.downsample_layers.{i}.1.weight', F'backbone.encoder.stages.{i}.downsampling_layer.1.weight')							)
               rename_keys.append((F'backbone.downsample_layers.{i}.1.bias', F'backbone.encoder.stages.{i}.downsampling_layer.1.bias')							)

          rename_keys.append((F'backbone.norm{i}.weight', F'backbone.hidden_states_norms.stage{i+1}.weight')							)
          rename_keys.append((F'backbone.norm{i}.bias', F'backbone.hidden_states_norms.stage{i+1}.bias')							)

     # decode head
     rename_keys.extend(
         [
             ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
             ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
             ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
             ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
         ]							)
     # fmt: on

     return rename_keys



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[Any]      ,    UpperCAmelCase_    :				int      ,    UpperCAmelCase_    :				int							)     ->		Optional[int]:
     __lowerCamelCase   :	List[str]								=		dct.pop(UpperCAmelCase_							)
     __lowerCamelCase   :	Any								=		val



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Any      ,    UpperCAmelCase_    :				Optional[Any]      ,    UpperCAmelCase_    :				List[Any]							)     ->		Optional[Any]:
     __lowerCamelCase   :	List[Any]								=		{
         'upernet-convnext-tiny': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth',
         'upernet-convnext-small': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth',
         'upernet-convnext-base': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth',
         'upernet-convnext-large': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth',
         'upernet-convnext-xlarge': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth',
     }
     __lowerCamelCase   :	List[Any]								=		model_name_to_url[model_name]
     __lowerCamelCase   :	Any								=		torch.hub.load_state_dict_from_url(UpperCAmelCase_      ,    map_location='cpu'							)['state_dict']

     __lowerCamelCase   :	Tuple								=		get_upernet_config(UpperCAmelCase_							)
     __lowerCamelCase   :	Optional[int]								=		UperNetForSemanticSegmentation(UpperCAmelCase_							)
     model.eval()

     # replace "bn" => "batch_norm"
     for key in state_dict.copy().keys():
          __lowerCamelCase   :	str								=		state_dict.pop(UpperCAmelCase_							)
          if "bn" in key:
               __lowerCamelCase   :	Union[str, Any]								=		key.replace('bn'      ,    'batch_norm'							)
          __lowerCamelCase   :	Union[str, Any]								=		val

     # rename keys
     __lowerCamelCase   :	Optional[int]								=		create_rename_keys(UpperCAmelCase_							)
     for src, dest in rename_keys:
          rename_key(UpperCAmelCase_      ,    UpperCAmelCase_      ,    UpperCAmelCase_							)

     model.load_state_dict(UpperCAmelCase_							)

     # verify on image
     __lowerCamelCase   :	Union[str, Any]								=		'https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'
     __lowerCamelCase   :	Optional[Any]								=		Image.open(requests.get(UpperCAmelCase_      ,    stream=UpperCAmelCase_							).raw							).convert('RGB'							)

     __lowerCamelCase   :	Tuple								=		SegformerImageProcessor()
     __lowerCamelCase   :	Tuple								=		processor(UpperCAmelCase_      ,    return_tensors='pt'							).pixel_values

     with torch.no_grad():
          __lowerCamelCase   :	Dict								=		model(UpperCAmelCase_							)

     if model_name == "upernet-convnext-tiny":
          __lowerCamelCase   :	Optional[int]								=		torch.tensor(
              [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]							)
     elif model_name == "upernet-convnext-small":
          __lowerCamelCase   :	List[str]								=		torch.tensor(
              [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]]							)
     elif model_name == "upernet-convnext-base":
          __lowerCamelCase   :	Dict								=		torch.tensor(
              [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]]							)
     elif model_name == "upernet-convnext-large":
          __lowerCamelCase   :	str								=		torch.tensor(
              [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]]							)
     elif model_name == "upernet-convnext-xlarge":
          __lowerCamelCase   :	Optional[int]								=		torch.tensor(
              [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]]							)
     print('Logits:'      ,    outputs.logits[0, 0, :3, :3]							)
     assert torch.allclose(outputs.logits[0, 0, :3, :3]      ,    UpperCAmelCase_      ,    atol=1e-4							)
     print('Looks ok!'							)

     if pytorch_dump_folder_path is not None:
          print(F'Saving model {model_name} to {pytorch_dump_folder_path}'							)
          model.save_pretrained(UpperCAmelCase_							)
          print(F'Saving processor to {pytorch_dump_folder_path}'							)
          processor.save_pretrained(UpperCAmelCase_							)

     if push_to_hub:
          print(F'Pushing model and processor for {model_name} to hub'							)
          model.push_to_hub(F'openmmlab/{model_name}'							)
          processor.push_to_hub(F'openmmlab/{model_name}'							)


if __name__ == "__main__":
      A__		:					List[str]												=  argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--model_name""",
          default="""upernet-convnext-tiny""",
          type=str,
          choices=[f'''upernet-convnext-{size}''' for size in ["""tiny""", """small""", """base""", """large""", """xlarge"""]],
          help="""Name of the ConvNext UperNet model you'd like to convert.""",
      )
      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."""
      )

      A__		:					str												=  parser.parse_args()
      convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''



import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


A__		:					List[str]												=  get_logger(__name__)


A__		:					str												=  R"""
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

"""





class    UpperCAmelCase_  :





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'	)







class    UpperCAmelCase_  :





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'	)







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         for processor in self:
              __lowerCamelCase   :	str								=		inspect.signature(processor.__call__	).parameters
              if len(SCREAMING_SNAKE_CASE_	) > 3:
                   if not all(arg in kwargs for arg in list(function_args.keys()	)[2:]	):
                        raise ValueError(
                            f'Make sure that all the required parameters: {list(function_args.keys()	)} for '
                            f'{processor.__class__} are passed to the logits processor.'	)
                   __lowerCamelCase   :	Tuple								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)
              else:
                   __lowerCamelCase   :	int								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or not (temperature > 0):
              raise ValueError(f'`temperature` has to be a strictly positive float, but is {temperature}'	)

         __lowerCamelCase   :	Optional[int]								=		temperature







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Dict								=		scores / self.temperature
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     Union[str, Any]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or (top_p < 0 or top_p > 1.0):
              raise ValueError(f'`top_p` has to be a float > 0 and < 1, but is {top_p}'	)
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or (min_tokens_to_keep < 1):
              raise ValueError(f'`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}'	)

         __lowerCamelCase   :	str								=		top_p
         __lowerCamelCase   :	Tuple								=		filter_value
         __lowerCamelCase   :	Tuple								=		min_tokens_to_keep







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	Any								=		lax.top_k(SCREAMING_SNAKE_CASE_			,     scores.shape[-1]	)

         __lowerCamelCase   :	int								=		jnp.full_like(SCREAMING_SNAKE_CASE_			,     self.filter_value	)
         __lowerCamelCase   :	Tuple								=		jax.nn.softmax(SCREAMING_SNAKE_CASE_			,     axis=-1	).cumsum(axis=-1	)
         __lowerCamelCase   :	List[str]								=		cumulative_probs < self.top_p

         # include the token that is higher than top_p as well
         __lowerCamelCase   :	Tuple								=		jnp.roll(SCREAMING_SNAKE_CASE_			,     1	)
         score_mask |= score_mask.at[:, 0].set(SCREAMING_SNAKE_CASE_	)

         # min tokens to keep
         __lowerCamelCase   :	Any								=		score_mask.at[:, : self.min_tokens_to_keep].set(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		jnp.where(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jax.lax.sort_key_val(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)[-1]

         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     str:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or top_k <= 0:
              raise ValueError(f'`top_k` has to be a strictly positive integer, but is {top_k}'	)

         __lowerCamelCase   :	List[str]								=		max(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		filter_value







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	List[Any]								=		scores.shape
         __lowerCamelCase   :	Tuple								=		jnp.full(batch_size * vocab_size			,     self.filter_value	)

         __lowerCamelCase   :	int								=		min(self.top_k			,     scores.shape[-1]	)  # Safety check
         __lowerCamelCase				,       __lowerCamelCase   :	Tuple								=		lax.top_k(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Union[str, Any]								=		jnp.broadcast_to((jnp.arange(SCREAMING_SNAKE_CASE_	) * vocab_size)[:, None]			,     (batch_size, topk)	).flatten()
         __lowerCamelCase   :	List[Any]								=		topk_scores.flatten()
         __lowerCamelCase   :	Union[str, Any]								=		topk_indices.flatten() + shift

         __lowerCamelCase   :	Tuple								=		next_scores_flat.at[topk_indices_flat].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		next_scores_flat.reshape(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Any								=		bos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Optional[Any]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Optional[Any]								=		1 - jnp.bool_(cur_len - 1	)

         __lowerCamelCase   :	List[Any]								=		jnp.where(SCREAMING_SNAKE_CASE_			,     new_scores.at[:, self.bos_token_id].set(0	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Tuple								=		max_length
         __lowerCamelCase   :	Any								=		eos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	List[str]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Any								=		1 - jnp.bool_(cur_len - self.max_length + 1	)

         __lowerCamelCase   :	List[str]								=		jnp.where(SCREAMING_SNAKE_CASE_			,     new_scores.at[:, self.eos_token_id].set(0	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or min_length < 0:
              raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}'	)

         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or eos_token_id < 0:
              raise ValueError(f'`eos_token_id` has to be a positive integer, but is {eos_token_id}'	)

         __lowerCamelCase   :	str								=		min_length
         __lowerCamelCase   :	Optional[int]								=		eos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         # create boolean flag to decide if min length penalty should be applied
         __lowerCamelCase   :	Optional[Any]								=		1 - jnp.clip(cur_len - self.min_length			,     0			,     1	)

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.eos_token_id].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		list(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		begin_index







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	List[Any]								=		1 - jnp.bool_(cur_len - self.begin_index	)

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.begin_suppress_tokens].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         __lowerCamelCase   :	Tuple								=		list(SCREAMING_SNAKE_CASE_	)







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	int								=		scores.at[..., self.suppress_tokens].set(-float('inf'	)	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Optional[int]								=		dict(SCREAMING_SNAKE_CASE_	)
         # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
         # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
         # Indexes without forced tokens will have a negative value.
         __lowerCamelCase   :	Dict								=		jnp.ones((max(force_token_map.keys()	) + 1)			,     dtype=jnp.intaa	) * -1
         for index, token in force_token_map.items():
              if token is not None:
                   __lowerCamelCase   :	str								=		force_token_array.at[index].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		jnp.intaa(SCREAMING_SNAKE_CASE_	)







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         def _force_token(SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	List[str]								=		scores.shape[0]
              __lowerCamelCase   :	Tuple								=		self.force_token_array[generation_idx]

              __lowerCamelCase   :	List[Any]								=		jnp.ones_like(SCREAMING_SNAKE_CASE_			,     dtype=scores.dtype	) * -float('inf'	)
              __lowerCamelCase   :	Any								=		jnp.zeros((batch_size, 1)			,     dtype=scores.dtype	)
              __lowerCamelCase   :	str								=		lax.dynamic_update_slice(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     (0, current_token)	)
              return new_scores

         __lowerCamelCase   :	int								=		lax.cond(
             cur_len >= self.force_token_array.shape[0]			,     lambda: scores			,     lambda: lax.cond(
                 self.force_token_array[cur_len] >= 0			,     lambda: _force_token(SCREAMING_SNAKE_CASE_	)			,     lambda: scores			,     )			,     )
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	Any								=		generate_config.eos_token_id
         __lowerCamelCase   :	Dict								=		generate_config.no_timestamps_token_id
         __lowerCamelCase   :	Tuple								=		generate_config.no_timestamps_token_id + 1

         __lowerCamelCase   :	List[str]								=		decoder_input_length + 1

         if generate_config.is_multilingual:
              # room for language token and task token
              self.begin_index += 2
         if hasattr(SCREAMING_SNAKE_CASE_			,     'max_initial_timestamp_index'	):
              __lowerCamelCase   :	str								=		generate_config.max_initial_timestamp_index
         else:
              __lowerCamelCase   :	Optional[int]								=		model_config.vocab_size
         if self.max_initial_timestamp_index is None:
              __lowerCamelCase   :	Tuple								=		model_config.vocab_size







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         # suppress <|notimestamps|> which is handled by without_timestamps
         __lowerCamelCase   :	Union[str, Any]								=		scores.at[:, self.no_timestamps_token_id].set(-float('inf'	)	)

         def handle_pairs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	str								=		jnp.where((cur_len - self.begin_index) >= 1			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	List[Any]								=		jnp.where(
                  input_ids_k[cur_len - 1] >= self.timestamp_begin			,     True and last_was_timestamp			,     SCREAMING_SNAKE_CASE_			,     )

              __lowerCamelCase   :	Any								=		jnp.where((cur_len - self.begin_index) < 2			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	str								=		jnp.where(
                  input_ids_k[cur_len - 2] >= self.timestamp_begin			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     )

              return jnp.where(
                  SCREAMING_SNAKE_CASE_			,     jnp.where(
                      penultimate_was_timestamp > 0			,     scores_k.at[self.timestamp_begin :].set(-float('inf'	)	)			,     scores_k.at[: self.eos_token_id].set(-float('inf'	)	)			,     )			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	List[Any]								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		jnp.where(cur_len == self.begin_index			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jnp.where(
             self.max_initial_timestamp_index is not None			,     True and apply_max_initial_timestamp			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	Any								=		self.timestamp_begin + self.max_initial_timestamp_index

         __lowerCamelCase   :	str								=		jnp.where(
             SCREAMING_SNAKE_CASE_			,     scores.at[:, last_allowed + 1 :].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_			,     )

         # if sum of probability over timestamps is above any other token, sample timestamp
         __lowerCamelCase   :	Any								=		jax.nn.log_softmax(SCREAMING_SNAKE_CASE_			,     axis=-1	)

         def handle_cumulative_probs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	Union[str, Any]								=		jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]			,     axis=-1	)
              __lowerCamelCase   :	List[str]								=		jnp.max(logprobs_k[: self.timestamp_begin]	)
              return jnp.where(
                  timestamp_logprob > max_text_token_logprob			,     scores_k.at[: self.timestamp_begin].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	Dict								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         return scores






