
"""simple docstring"""




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 							UpperCAmelCase				(  unittest.TestCase  ):




	'''simple docstring'''





	def     UpperCamelCase_  (	self :			List[str]							,A :			Dict							,A :			Tuple			):
					__A             =			jnp.ones((batch_size, length)			) / length
					return scores






	def     UpperCamelCase_  (	self :			Union[str, Any]			):
					__A             =			None
					__A             =			20

					__A             =			self._get_uniform_logits(batch_size=2							,length=SCREAMING_SNAKE_CASE_			)

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

					# compute softmax
					__A             =			jax.nn.softmax(SCREAMING_SNAKE_CASE_							,axis=-1			)

					__A             =			FlaxTemperatureLogitsWarper(temperature=0.5			)
					__A             =			FlaxTemperatureLogitsWarper(temperature=1.3			)

					__A             =			jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE_							,scores.copy()							,cur_len=SCREAMING_SNAKE_CASE_			)							,axis=-1			)
					__A             =			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     UpperCamelCase_  (	self :			List[Any]			):
					__A             =			None
					__A             =			10
					__A             =			2

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

					__A             =			FlaxTopKLogitsWarper(3			)

					__A             =			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             =			5
					__A             =			FlaxTopKLogitsWarper(top_k=1							,filter_value=0.0							,min_tokens_to_keep=3			)

					__A             =			np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_			)[None, :]							,(batch_size, length)			).copy()
					__A             =			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     UpperCamelCase_  (	self :			Dict			):
					__A             =			None
					__A             =			10
					__A             =			2

					# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
					__A             =			np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]			)			)

					__A             =			FlaxTopPLogitsWarper(0.8			)
					__A             =			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             =			np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]			)
					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,atol=1E-3			)			)

					# check edge cases with negative and extreme logits
					__A             =			np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_			)[None, :]							,(batch_size, vocab_size)			).copy() - (
					    vocab_size // 2
					)

					# make ramp_logits more extreme
					__A             =			ramp_logits[1] * 1_00.0

					# make sure at least 2 tokens are kept
					__A             =			FlaxTopPLogitsWarper(0.9							,min_tokens_to_keep=2							,filter_value=0.0			)
					__A             =			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     UpperCamelCase_  (	self :			Optional[int]			):
					__A             =			20
					__A             =			4
					__A             =			0

					__A             =			FlaxMinLengthLogitsProcessor(min_length=10							,eos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that min length is applied at length 5
					__A             =			ids_tensor((batch_size, 20)							,vocab_size=20			)
					__A             =			5
					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			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             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			15
					__A             =			min_dist_processor(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)






	def     UpperCamelCase_  (	self :			Dict			):
					__A             =			20
					__A             =			4
					__A             =			0

					__A             =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the bos_token_id score
					__A             =			ids_tensor((batch_size, 1)							,vocab_size=20			)
					__A             =			1
					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			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             =			3
					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			logits_processor(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)






	def     UpperCamelCase_  (	self :			Union[str, Any]			):
					__A             =			20
					__A             =			4
					__A             =			0
					__A             =			5

					__A             =			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             =			ids_tensor((batch_size, 4)							,vocab_size=20			)
					__A             =			4
					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			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             =			3
					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			logits_processor(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)






	def     UpperCamelCase_  (	self :			str			):
					__A             =			4
					__A             =			10
					__A             =			15
					__A             =			2
					__A             =			1
					__A             =			15

					# dummy input_ids and scores
					__A             =			ids_tensor((batch_size, sequence_length)							,SCREAMING_SNAKE_CASE_			)
					__A             =			input_ids.copy()

					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			scores.copy()

					# instantiate all dist processors
					__A             =			FlaxTemperatureLogitsWarper(temperature=0.5			)
					__A             =			FlaxTopKLogitsWarper(3			)
					__A             =			FlaxTopPLogitsWarper(0.8			)

					# instantiate all logits processors
					__A             =			FlaxMinLengthLogitsProcessor(min_length=10							,eos_token_id=SCREAMING_SNAKE_CASE_			)
					__A             =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)
					__A             =			FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_							,eos_token_id=SCREAMING_SNAKE_CASE_			)

					__A             =			10

					# no processor list
					__A             =			temp_dist_warp(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					__A             =			top_k_warp(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					__A             =			top_p_warp(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					__A             =			min_dist_proc(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					__A             =			bos_dist_proc(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
					__A             =			eos_dist_proc(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)

					# with processor list
					__A             =			FlaxLogitsProcessorList(
					    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
					__A             =			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     UpperCamelCase_  (	self :			Dict			):
					__A             =			4
					__A             =			10
					__A             =			15
					__A             =			2
					__A             =			1
					__A             =			15

					# dummy input_ids and scores
					__A             =			ids_tensor((batch_size, sequence_length)							,SCREAMING_SNAKE_CASE_			)
					__A             =			input_ids.copy()

					__A             =			self._get_uniform_logits(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			scores.copy()

					# instantiate all dist processors
					__A             =			FlaxTemperatureLogitsWarper(temperature=0.5			)
					__A             =			FlaxTopKLogitsWarper(3			)
					__A             =			FlaxTopPLogitsWarper(0.8			)

					# instantiate all logits processors
					__A             =			FlaxMinLengthLogitsProcessor(min_length=10							,eos_token_id=SCREAMING_SNAKE_CASE_			)
					__A             =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)
					__A             =			FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_							,eos_token_id=SCREAMING_SNAKE_CASE_			)

					__A             =			10

					# no processor list
					def run_no_processor_list(A :			Optional[int]							,A :			List[Any]							,A :			Tuple			):
									__A             =			temp_dist_warp(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									__A             =			top_k_warp(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									__A             =			top_p_warp(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									__A             =			min_dist_proc(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									__A             =			bos_dist_proc(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									__A             =			eos_dist_proc(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					# with processor list
					def run_processor_list(A :			Any							,A :			List[Any]							,A :			Optional[int]			):
									__A             =			FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
									__A             =			processor(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					__A             =			jax.jit(SCREAMING_SNAKE_CASE_			)
					__A             =			jax.jit(SCREAMING_SNAKE_CASE_			)

					__A             =			jitted_run_no_processor_list(SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_							,SCREAMING_SNAKE_CASE_			)
					__A             =			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()			)


from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin





class 							UpperCAmelCase				(  __SCREAMING_SNAKE_CASE      ,      __SCREAMING_SNAKE_CASE      ,      __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	snake_case_    =					[R"h\.\d+\.attn\.bias", R"h\.\d+\.attn\.masked_bias"]






	@register_to_config
	def __init__(	self :			Tuple							,A :			int							,A :			int							,A :			Optional[int] = None							,A :			int = 5_02_57							,A :			int = 10_24							,A :			int = 7_68							,A :			int = 12							,A :			int = 12							,A :			Optional[int] = None							,A :			str = "gelu_new"							,A :			float = 0.1							,A :			float = 0.1							,A :			float = 0.1							,A :			float = 1E-5							,A :			float = 0.02							,A :			bool = True							,A :			bool = True							,A :			bool = False							,A :			bool = False							,):
					super().__init__()

					__A             =			prefix_length

					if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
									raise ValueError(
									    f'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
									    f''' `n_embd`: {n_embd} are not equal.'''			)

					__A             =			prefix_inner_dim
					__A             =			prefix_hidden_dim

					__A             =			(
					    nn.Linear(self.prefix_inner_dim							,self.prefix_hidden_dim			)
					    if self.prefix_hidden_dim is not None
					    else nn.Identity()
					)
					__A             =			(
					    nn.Linear(self.prefix_hidden_dim							,A			) if self.prefix_hidden_dim is not None else nn.Identity()
					)

					__A             =			GPTaConfig(
					    vocab_size=A							,n_positions=A							,n_embd=A							,n_layer=A							,n_head=A							,n_inner=A							,activation_function=A							,resid_pdrop=A							,embd_pdrop=A							,attn_pdrop=A							,layer_norm_epsilon=A							,initializer_range=A							,scale_attn_weights=A							,use_cache=A							,scale_attn_by_inverse_layer_idx=A							,reorder_and_upcast_attn=A							,)
					__A             =			GPTaLMHeadModel(A			)






	def     UpperCamelCase_  (	self :			Union[str, Any]							,A :			torch.Tensor							,A :			torch.Tensor							,A :			Optional[torch.Tensor] = None							,A :			Optional[torch.Tensor] = None							,):
					__A             =			self.transformer.transformer.wte(A			)
					__A             =			self.encode_prefix(A			)
					__A             =			self.decode_prefix(A			)
					__A             =			torch.cat((prefix_embeds, embedding_text)							,dim=1			)

					if labels is not None:
									__A             =			self.get_dummy_token(input_ids.shape[0]							,input_ids.device			)
									__A             =			torch.cat((dummy_token, input_ids)							,dim=1			)
					__A             =			self.transformer(inputs_embeds=A							,labels=A							,attention_mask=A			)
					if self.prefix_hidden_dim is not None:
									return out, hidden
					else:
									return out






	def     UpperCamelCase_  (	self :			int							,A :			int							,A :			torch.device			):
					return torch.zeros(A							,self.prefix_length							,dtype=torch.intaa							,device=A			)






	def     UpperCamelCase_  (	self :			Optional[Any]							,A :			List[Any]			):
					return self.encode_prefix(A			)






	@torch.no_grad()
	def     UpperCamelCase_  (	self :			str							,A :			List[str]							,A :			List[Any]							,A :			List[Any]			):
					__A             =			torch.split(A							,1							,dim=0			)
					__A             =			[]
					__A             =			[]
					for feature in features:
									__A             =			self.decode_prefix(feature.to(A			)			)  # back to the clip feature
									# Only support beam search for now
									__A      ,							__A             =			self.generate_beam(
									    input_embeds=A							,device=A							,eos_token_id=A			)
									generated_tokens.append(output_tokens[0]			)
									generated_seq_lengths.append(seq_lengths[0]			)
					__A             =			torch.stack(A			)
					__A             =			torch.stack(A			)
					return generated_tokens, generated_seq_lengths







	@torch.no_grad()
	def     UpperCamelCase_  (	self :			Tuple							,A :			List[str]=None							,A :			Dict=None							,A :			Optional[int]=None							,A :			int = 5							,A :			int = 67							,A :			float = 1.0							,A :			Optional[int] = None							,):
					__A             =			eos_token_id
					__A             =			None
					__A             =			None
					__A             =			torch.ones(A							,device=A							,dtype=torch.int			)
					__A             =			torch.zeros(A							,device=A							,dtype=torch.bool			)

					if input_embeds is not None:
									__A             =			input_embeds
					else:
									__A             =			self.transformer.transformer.wte(A			)

					for i in range(A			):
									__A             =			self.transformer(inputs_embeds=A			)
									__A             =			outputs.logits
									__A             =			logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
									__A             =			logits.softmax(-1			).log()

									if scores is None:
													__A      ,							__A             =			logits.topk(A							,-1			)
													__A             =			generated.expand(A							,*generated.shape[1:]			)
													__A      ,							__A             =			next_tokens.permute(1							,0			), scores.squeeze(0			)
													if tokens is None:
																	__A             =			next_tokens
													else:
																	__A             =			tokens.expand(A							,*tokens.shape[1:]			)
																	__A             =			torch.cat((tokens, next_tokens)							,dim=1			)
									else:
													__A             =			-float(np.inf			)
													__A             =			0
													__A             =			scores[:, None] + logits
													seq_lengths[~is_stopped] += 1
													__A             =			scores_sum / seq_lengths[:, None]
													__A      ,							__A             =			scores_sum_average.view(-1			).topk(A							,-1			)
													__A             =			next_tokens // scores_sum.shape[1]
													__A             =			seq_lengths[next_tokens_source]
													__A             =			next_tokens % scores_sum.shape[1]
													__A             =			next_tokens.unsqueeze(1			)
													__A             =			tokens[next_tokens_source]
													__A             =			torch.cat((tokens, next_tokens)							,dim=1			)
													__A             =			generated[next_tokens_source]
													__A             =			scores_sum_average * seq_lengths
													__A             =			is_stopped[next_tokens_source]

									__A             =			self.transformer.transformer.wte(next_tokens.squeeze()			).view(generated.shape[0]							,1							,-1			)
									__A             =			torch.cat((generated, next_token_embed)							,dim=1			)
									__A             =			is_stopped + next_tokens.eq(A			).squeeze()
									if is_stopped.all():
													break

					__A             =			scores / seq_lengths
					__A             =			scores.argsort(descending=A			)
					# tokens tensors are already padded to max_seq_length
					__A             =			[tokens[i] for i in order]
					__A             =			torch.stack(A							,dim=0			)
					__A             =			torch.tensor([seq_lengths[i] for i in order]							,dtype=seq_lengths.dtype			)
					return output_texts, seq_lengths
