
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      lowerCamelCase__       (       __lowerCAmelCase						:						Dict       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			384
			if "tiny" in model_name:
						lowerCAmelCase_      			=			[3, 3, 9, 3]
						lowerCAmelCase_      			=			[96, 192, 384, 768]
			if "small" in model_name:
						lowerCAmelCase_      			=			[3, 3, 27, 3]
						lowerCAmelCase_      			=			[96, 192, 384, 768]
			if "base" in model_name:
						lowerCAmelCase_      			=			[3, 3, 27, 3]
						lowerCAmelCase_      			=			[128, 256, 512, 1024]
						lowerCAmelCase_      			=			512
			if "large" in model_name:
						lowerCAmelCase_      			=			[3, 3, 27, 3]
						lowerCAmelCase_      			=			[192, 384, 768, 1536]
						lowerCAmelCase_      			=			768
			if "xlarge" in model_name:
						lowerCAmelCase_      			=			[3, 3, 27, 3]
						lowerCAmelCase_      			=			[256, 512, 1024, 2048]
						lowerCAmelCase_      			=			1024

			# set label information
			lowerCAmelCase_      			=			150
			lowerCAmelCase_      			=			"huggingface/label-files"
			lowerCAmelCase_      			=			"ade20k-id2label.json"
			lowerCAmelCase_      			=			json.load(open(hf_hub_download(__lowerCAmelCase					,							__lowerCAmelCase					,							repo_type="dataset"       )					,							"r"       )       )
			lowerCAmelCase_      			=			{int(__lowerCAmelCase       ): v for k, v in idalabel.items()}
			lowerCAmelCase_      			=			{v: k for k, v in idalabel.items()}

			lowerCAmelCase_      			=			ConvNextConfig(
			    depths=__lowerCAmelCase					,							hidden_sizes=__lowerCAmelCase					,							out_features=["stage1", "stage2", "stage3", "stage4"]       )
			lowerCAmelCase_      			=			UperNetConfig(
			    backbone_config=__lowerCAmelCase					,							auxiliary_in_channels=__lowerCAmelCase					,							num_labels=__lowerCAmelCase					,							idalabel=__lowerCAmelCase					,							labelaid=__lowerCAmelCase					,							)

			return config






def      lowerCamelCase__       (       __lowerCAmelCase						:						Dict       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			[]

			# 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      lowerCamelCase__       (       __lowerCAmelCase						:						Dict					,							__lowerCAmelCase						:						List[Any]					,							__lowerCAmelCase						:						List[str]       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			dct.pop(__lowerCAmelCase       )
			lowerCAmelCase_      			=			val






def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[int]					,							__lowerCAmelCase						:						str					,							__lowerCAmelCase						:						Optional[Any]       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			{
			    "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_      			=			model_name_to_url[model_name]
			lowerCAmelCase_      			=			torch.hub.load_state_dict_from_url(__lowerCAmelCase					,							map_location="cpu"       )["state_dict"]

			lowerCAmelCase_      			=			get_upernet_config(__lowerCAmelCase       )
			lowerCAmelCase_      			=			UperNetForSemanticSegmentation(__lowerCAmelCase       )
			model.eval()

			# replace "bn" => "batch_norm"
			for key in state_dict.copy().keys():
						lowerCAmelCase_      			=			state_dict.pop(__lowerCAmelCase       )
						if "bn" in key:
									lowerCAmelCase_      			=			key.replace("bn"					,							"batch_norm"       )
						lowerCAmelCase_      			=			val

			# rename keys
			lowerCAmelCase_      			=			create_rename_keys(__lowerCAmelCase       )
			for src, dest in rename_keys:
						rename_key(__lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase       )

			model.load_state_dict(__lowerCAmelCase       )

			# verify on image
			lowerCAmelCase_      			=			"https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
			lowerCAmelCase_      			=			Image.open(requests.get(__lowerCAmelCase					,							stream=__lowerCAmelCase       ).raw       ).convert("RGB"       )

			lowerCAmelCase_      			=			SegformerImageProcessor()
			lowerCAmelCase_      			=			processor(__lowerCAmelCase					,							return_tensors="pt"       ).pixel_values

			with torch.no_grad():
						lowerCAmelCase_      			=			model(__lowerCAmelCase       )

			if model_name == "upernet-convnext-tiny":
						lowerCAmelCase_      			=			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_      			=			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_      			=			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_      			=			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_      			=			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]					,							__lowerCAmelCase					,							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(__lowerCAmelCase       )
						print(F"""Saving processor to {pytorch_dump_folder_path}"""       )
						processor.save_pretrained(__lowerCAmelCase       )

			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     =  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     =  parser.parse_args()
			convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





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     =  get_logger(__name__)


_A     =  R"\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n"





class 			_lowerCAmelCase      :




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





class 			_lowerCAmelCase      :




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





class 			_lowerCAmelCase      (			__a ):




				@add_start_docstrings(_UpperCamelCase  )
				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       **_UpperCamelCase  )       ->   jnp.ndarray:
							for processor in self:
										lowerCAmelCase_      			=			inspect.signature(processor.__call__  ).parameters
										if len(_UpperCamelCase  ) > 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_      			=			processor(_UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       **_UpperCamelCase  )
										else:
													lowerCAmelCase_      			=			processor(_UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )
							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase  )       ->   Tuple:
							if not isinstance(_UpperCamelCase       ,       _UpperCamelCase  ) or not (temperature > 0):
										raise ValueError(f"""`temperature` has to be a strictly positive float, but is {temperature}"""  )

							lowerCAmelCase_      			=			temperature




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							lowerCAmelCase_      			=			scores / self.temperature
							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase = -float("Inf"  )       ,       _UpperCamelCase = 1  )       ->   Union[str, Any]:
							if not isinstance(_UpperCamelCase       ,       _UpperCamelCase  ) 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(_UpperCamelCase       ,       _UpperCamelCase  ) 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_      			=			top_p
							lowerCAmelCase_      			=			filter_value
							lowerCAmelCase_      			=			min_tokens_to_keep




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

							lowerCAmelCase_      			=			jnp.full_like(_UpperCamelCase       ,       self.filter_value  )
							lowerCAmelCase_      			=			jax.nn.softmax(_UpperCamelCase       ,       axis=-1  ).cumsum(axis=-1  )
							lowerCAmelCase_      			=			cumulative_probs < self.top_p

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

							# min tokens to keep
							lowerCAmelCase_      			=			score_mask.at[:, : self.min_tokens_to_keep].set(_UpperCamelCase  )

							lowerCAmelCase_      			=			jnp.where(_UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			jax.lax.sort_key_val(_UpperCamelCase       ,       _UpperCamelCase  )[-1]

							return next_scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase = -float("Inf"  )       ,       _UpperCamelCase = 1  )       ->   List[Any]:
							if not isinstance(_UpperCamelCase       ,       _UpperCamelCase  ) or top_k <= 0:
										raise ValueError(f"""`top_k` has to be a strictly positive integer, but is {top_k}"""  )

							lowerCAmelCase_      			=			max(_UpperCamelCase       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			filter_value




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							lowerCAmelCase_			,						lowerCAmelCase_      			=			scores.shape
							lowerCAmelCase_      			=			jnp.full(batch_size * vocab_size       ,       self.filter_value  )

							lowerCAmelCase_      			=			min(self.top_k       ,       scores.shape[-1]  )  # Safety check
							lowerCAmelCase_			,						lowerCAmelCase_      			=			lax.top_k(_UpperCamelCase       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			jnp.broadcast_to((jnp.arange(_UpperCamelCase  ) * vocab_size)[:, None]       ,       (batch_size, topk)  ).flatten()
							lowerCAmelCase_      			=			topk_scores.flatten()
							lowerCAmelCase_      			=			topk_indices.flatten() + shift

							lowerCAmelCase_      			=			next_scores_flat.at[topk_indices_flat].set(_UpperCamelCase  )
							lowerCAmelCase_      			=			next_scores_flat.reshape(_UpperCamelCase       ,       _UpperCamelCase  )
							return next_scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase  )       ->   Any:
							lowerCAmelCase_      			=			bos_token_id




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							lowerCAmelCase_      			=			jnp.full(scores.shape       ,       -float("inf"  )  )

							lowerCAmelCase_      			=			1 - jnp.bool_(cur_len - 1  )

							lowerCAmelCase_      			=			jnp.where(_UpperCamelCase       ,       new_scores.at[:, self.bos_token_id].set(0  )       ,       _UpperCamelCase  )

							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   List[str]:
							lowerCAmelCase_      			=			max_length
							lowerCAmelCase_      			=			eos_token_id




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							lowerCAmelCase_      			=			jnp.full(scores.shape       ,       -float("inf"  )  )

							lowerCAmelCase_      			=			1 - jnp.bool_(cur_len - self.max_length + 1  )

							lowerCAmelCase_      			=			jnp.where(_UpperCamelCase       ,       new_scores.at[:, self.eos_token_id].set(0  )       ,       _UpperCamelCase  )

							return scores





class 			_lowerCAmelCase      (			__a ):




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

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

							lowerCAmelCase_      			=			min_length
							lowerCAmelCase_      			=			eos_token_id




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							# create boolean flag to decide if min length penalty should be applied
							lowerCAmelCase_      			=			1 - jnp.clip(cur_len - self.min_length       ,       0       ,       1  )

							lowerCAmelCase_      			=			jnp.where(_UpperCamelCase       ,       scores.at[:, self.eos_token_id].set(-float("inf"  )  )       ,       _UpperCamelCase  )

							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   int:
							lowerCAmelCase_      			=			list(_UpperCamelCase  )
							lowerCAmelCase_      			=			begin_index




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   int:
							lowerCAmelCase_      			=			1 - jnp.bool_(cur_len - self.begin_index  )

							lowerCAmelCase_      			=			jnp.where(_UpperCamelCase       ,       scores.at[:, self.begin_suppress_tokens].set(-float("inf"  )  )       ,       _UpperCamelCase  )

							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase  )       ->   Optional[Any]:
							lowerCAmelCase_      			=			list(_UpperCamelCase  )




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							lowerCAmelCase_      			=			scores.at[..., self.suppress_tokens].set(-float("inf"  )  )

							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase  )       ->   List[Any]:
							lowerCAmelCase_      			=			dict(_UpperCamelCase  )
							# 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_      			=			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_      			=			force_token_array.at[index].set(_UpperCamelCase  )
							lowerCAmelCase_      			=			jnp.intaa(_UpperCamelCase  )




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   jnp.ndarray:
							def _force_token(_UpperCamelCase  ):
										lowerCAmelCase_      			=			scores.shape[0]
										lowerCAmelCase_      			=			self.force_token_array[generation_idx]

										lowerCAmelCase_      			=			jnp.ones_like(_UpperCamelCase       ,       dtype=scores.dtype  ) * -float("inf"  )
										lowerCAmelCase_      			=			jnp.zeros((batch_size, 1)       ,       dtype=scores.dtype  )
										lowerCAmelCase_      			=			lax.dynamic_update_slice(_UpperCamelCase       ,       _UpperCamelCase       ,       (0, current_token)  )
										return new_scores

							lowerCAmelCase_      			=			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(_UpperCamelCase  )       ,       lambda: scores       ,       )       ,       )
							return scores





class 			_lowerCAmelCase      (			__a ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Tuple:
							lowerCAmelCase_      			=			generate_config.eos_token_id
							lowerCAmelCase_      			=			generate_config.no_timestamps_token_id
							lowerCAmelCase_      			=			generate_config.no_timestamps_token_id + 1

							lowerCAmelCase_      			=			decoder_input_length + 1

							if generate_config.is_multilingual:
										# room for language token and task token
										self.begin_index += 2
							if hasattr(_UpperCamelCase       ,       "max_initial_timestamp_index"  ):
										lowerCAmelCase_      			=			generate_config.max_initial_timestamp_index
							else:
										lowerCAmelCase_      			=			model_config.vocab_size
							if self.max_initial_timestamp_index is None:
										lowerCAmelCase_      			=			model_config.vocab_size




				def __call__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Optional[Any]:
							# suppress <|notimestamps|> which is handled by without_timestamps
							lowerCAmelCase_      			=			scores.at[:, self.no_timestamps_token_id].set(-float("inf"  )  )

							def handle_pairs(_UpperCamelCase       ,       _UpperCamelCase  ):
										lowerCAmelCase_      			=			jnp.where((cur_len - self.begin_index) >= 1       ,       _UpperCamelCase       ,       _UpperCamelCase  )
										lowerCAmelCase_      			=			jnp.where(
										    input_ids_k[cur_len - 1] >= self.timestamp_begin       ,       True and last_was_timestamp       ,       _UpperCamelCase       ,       )

										lowerCAmelCase_      			=			jnp.where((cur_len - self.begin_index) < 2       ,       _UpperCamelCase       ,       _UpperCamelCase  )
										lowerCAmelCase_      			=			jnp.where(
										    input_ids_k[cur_len - 2] >= self.timestamp_begin       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       )

										return jnp.where(
										    _UpperCamelCase       ,       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"  )  )       ,       )       ,       _UpperCamelCase       ,       )

							lowerCAmelCase_      			=			jax.vmap(_UpperCamelCase  )(_UpperCamelCase       ,       _UpperCamelCase  )

							lowerCAmelCase_      			=			jnp.where(cur_len == self.begin_index       ,       _UpperCamelCase       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			jnp.where(
							    self.max_initial_timestamp_index is not None       ,       True and apply_max_initial_timestamp       ,       _UpperCamelCase       ,       )

							lowerCAmelCase_      			=			self.timestamp_begin + self.max_initial_timestamp_index

							lowerCAmelCase_      			=			jnp.where(
							    _UpperCamelCase       ,       scores.at[:, last_allowed + 1 :].set(-float("inf"  )  )       ,       _UpperCamelCase       ,       )

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

							def handle_cumulative_probs(_UpperCamelCase       ,       _UpperCamelCase  ):
										lowerCAmelCase_      			=			jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]       ,       axis=-1  )
										lowerCAmelCase_      			=			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"  )  )       ,       _UpperCamelCase       ,       )

							lowerCAmelCase_      			=			jax.vmap(_UpperCamelCase  )(_UpperCamelCase       ,       _UpperCamelCase  )

							return scores



