






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 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[str]:
							UpperCAmelCase_    :      Tuple            =						jnp.ones((batch_size, length)   ) / length
							return scores



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							UpperCAmelCase_    :      List[str]            =						None
							UpperCAmelCase_    :      str            =						2_0

							UpperCAmelCase_    :      Tuple            =						self._get_uniform_logits(batch_size=2				,		length=__snake_case   )

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

							# compute softmax
							UpperCAmelCase_    :      Dict            =						jax.nn.softmax(__snake_case				,		axis=-1   )

							UpperCAmelCase_    :      Optional[Any]            =						FlaxTemperatureLogitsWarper(temperature=0.5   )
							UpperCAmelCase_    :      List[Any]            =						FlaxTemperatureLogitsWarper(temperature=1.3   )

							UpperCAmelCase_    :      Any            =						jax.nn.softmax(temp_dist_warper_sharper(__snake_case				,		scores.copy()				,		cur_len=__snake_case   )				,		axis=-1   )
							UpperCAmelCase_    :      str            =						jax.nn.softmax(temp_dist_warper_smoother(__snake_case				,		scores.copy()				,		cur_len=__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]:
							UpperCAmelCase_    :      List[str]            =						None
							UpperCAmelCase_    :      Union[str, Any]            =						1_0
							UpperCAmelCase_    :      Union[str, Any]            =						2

							# create ramp distribution
							UpperCAmelCase_    :      Union[str, Any]            =						np.broadcast_to(np.arange(__snake_case   )[None, :]				,		(batch_size, vocab_size)   ).copy()
							UpperCAmelCase_    :      Dict            =						ramp_logits[1:, : vocab_size // 2] + vocab_size

							UpperCAmelCase_    :      Optional[int]            =						FlaxTopKLogitsWarper(3   )

							UpperCAmelCase_    :      Optional[Any]            =						top_k_warp(__snake_case				,		__snake_case				,		cur_len=__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
							UpperCAmelCase_    :      Tuple            =						5
							UpperCAmelCase_    :      int            =						FlaxTopKLogitsWarper(top_k=1				,		filter_value=0.0				,		min_tokens_to_keep=3   )

							UpperCAmelCase_    :      Any            =						np.broadcast_to(np.arange(__snake_case   )[None, :]				,		(batch_size, length)   ).copy()
							UpperCAmelCase_    :      Tuple            =						top_k_warp_safety_check(__snake_case				,		__snake_case				,		cur_len=__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   )	->							Tuple:
							UpperCAmelCase_    :      Dict            =						None
							UpperCAmelCase_    :      Optional[Any]            =						1_0
							UpperCAmelCase_    :      Dict            =						2

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

							UpperCAmelCase_    :      Optional[int]            =						FlaxTopPLogitsWarper(0.8   )
							UpperCAmelCase_    :      Optional[int]            =						np.exp(top_p_warp(__snake_case				,		__snake_case				,		cur_len=__snake_case   )   )

							# dist should be filtered to keep min num values so that sum is >= top_p
							# exp (-inf) => 0
							UpperCAmelCase_    :      Optional[Any]            =						np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]   )
							self.assertTrue(np.allclose(__snake_case				,		__snake_case				,		atol=1E-3   )   )

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

							# make ramp_logits more extreme
							UpperCAmelCase_    :      Any            =						ramp_logits[1] * 1_00.0

							# make sure at least 2 tokens are kept
							UpperCAmelCase_    :      int            =						FlaxTopPLogitsWarper(0.9				,		min_tokens_to_keep=2				,		filter_value=0.0   )
							UpperCAmelCase_    :      Union[str, Any]            =						top_p_warp(__snake_case				,		__snake_case				,		cur_len=__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[Any]:
							UpperCAmelCase_    :      Any            =						2_0
							UpperCAmelCase_    :      str            =						4
							UpperCAmelCase_    :      int            =						0

							UpperCAmelCase_    :      Union[str, Any]            =						FlaxMinLengthLogitsProcessor(min_length=1_0				,		eos_token_id=__snake_case   )

							# check that min length is applied at length 5
							UpperCAmelCase_    :      Any            =						ids_tensor((batch_size, 2_0)				,		vocab_size=2_0   )
							UpperCAmelCase_    :      Tuple            =						5
							UpperCAmelCase_    :      List[str]            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      Optional[int]            =						min_dist_processor(__snake_case				,		__snake_case				,		cur_len=__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
							UpperCAmelCase_    :      Tuple            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      List[Any]            =						1_5
							UpperCAmelCase_    :      List[str]            =						min_dist_processor(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
							self.assertFalse(jnp.isinf(__snake_case   ).any()   )



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							UpperCAmelCase_    :      Tuple            =						2_0
							UpperCAmelCase_    :      Optional[int]            =						4
							UpperCAmelCase_    :      str            =						0

							UpperCAmelCase_    :      Tuple            =						FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case   )

							# check that all scores are -inf except the bos_token_id score
							UpperCAmelCase_    :      int            =						ids_tensor((batch_size, 1)				,		vocab_size=2_0   )
							UpperCAmelCase_    :      Any            =						1
							UpperCAmelCase_    :      Optional[int]            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      str            =						logits_processor(__snake_case				,		__snake_case				,		cur_len=__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
							UpperCAmelCase_    :      Any            =						3
							UpperCAmelCase_    :      Tuple            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      int            =						logits_processor(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
							self.assertFalse(jnp.isinf(__snake_case   ).any()   )



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_    :      List[str]            =						2_0
							UpperCAmelCase_    :      Union[str, Any]            =						4
							UpperCAmelCase_    :      Dict            =						0
							UpperCAmelCase_    :      Tuple            =						5

							UpperCAmelCase_    :      Union[str, Any]            =						FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case				,		eos_token_id=__snake_case   )

							# check that all scores are -inf except the eos_token_id when max_length is reached
							UpperCAmelCase_    :      int            =						ids_tensor((batch_size, 4)				,		vocab_size=2_0   )
							UpperCAmelCase_    :      Optional[int]            =						4
							UpperCAmelCase_    :      Optional[Any]            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      Optional[int]            =						logits_processor(__snake_case				,		__snake_case				,		cur_len=__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
							UpperCAmelCase_    :      Union[str, Any]            =						3
							UpperCAmelCase_    :      List[Any]            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      Dict            =						logits_processor(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
							self.assertFalse(jnp.isinf(__snake_case   ).any()   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      str            =						4
							UpperCAmelCase_    :      Dict            =						1_0
							UpperCAmelCase_    :      str            =						1_5
							UpperCAmelCase_    :      List[Any]            =						2
							UpperCAmelCase_    :      Optional[int]            =						1
							UpperCAmelCase_    :      Tuple            =						1_5

							# dummy input_ids and scores
							UpperCAmelCase_    :      str            =						ids_tensor((batch_size, sequence_length)				,		__snake_case   )
							UpperCAmelCase_    :      Any            =						input_ids.copy()

							UpperCAmelCase_    :      List[str]            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      str            =						scores.copy()

							# instantiate all dist processors
							UpperCAmelCase_    :      str            =						FlaxTemperatureLogitsWarper(temperature=0.5   )
							UpperCAmelCase_    :      int            =						FlaxTopKLogitsWarper(3   )
							UpperCAmelCase_    :      int            =						FlaxTopPLogitsWarper(0.8   )

							# instantiate all logits processors
							UpperCAmelCase_    :      Tuple            =						FlaxMinLengthLogitsProcessor(min_length=1_0				,		eos_token_id=__snake_case   )
							UpperCAmelCase_    :      Union[str, Any]            =						FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case   )
							UpperCAmelCase_    :      int            =						FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case				,		eos_token_id=__snake_case   )

							UpperCAmelCase_    :      Dict            =						1_0

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

							# with processor list
							UpperCAmelCase_    :      int            =						FlaxLogitsProcessorList(
							    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
							UpperCAmelCase_    :      List[Any]            =						processor(__snake_case				,		__snake_case				,		cur_len=__snake_case   )

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

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



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_    :      int            =						4
							UpperCAmelCase_    :      Tuple            =						1_0
							UpperCAmelCase_    :      Union[str, Any]            =						1_5
							UpperCAmelCase_    :      str            =						2
							UpperCAmelCase_    :      Optional[int]            =						1
							UpperCAmelCase_    :      Tuple            =						1_5

							# dummy input_ids and scores
							UpperCAmelCase_    :      Any            =						ids_tensor((batch_size, sequence_length)				,		__snake_case   )
							UpperCAmelCase_    :      Dict            =						input_ids.copy()

							UpperCAmelCase_    :      Optional[Any]            =						self._get_uniform_logits(__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      Optional[int]            =						scores.copy()

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

							# instantiate all logits processors
							UpperCAmelCase_    :      List[str]            =						FlaxMinLengthLogitsProcessor(min_length=1_0				,		eos_token_id=__snake_case   )
							UpperCAmelCase_    :      int            =						FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case   )
							UpperCAmelCase_    :      Union[str, Any]            =						FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case				,		eos_token_id=__snake_case   )

							UpperCAmelCase_    :      Optional[Any]            =						1_0

							# no processor list
							def run_no_processor_list(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      List[Any]            =						temp_dist_warp(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												UpperCAmelCase_    :      Optional[int]            =						top_k_warp(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												UpperCAmelCase_    :      Tuple            =						top_p_warp(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												UpperCAmelCase_    :      Union[str, Any]            =						min_dist_proc(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												UpperCAmelCase_    :      List[str]            =						bos_dist_proc(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												UpperCAmelCase_    :      str            =						eos_dist_proc(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												return scores

							# with processor list
							def run_processor_list(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      List[str]            =						FlaxLogitsProcessorList(
												    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
												UpperCAmelCase_    :      List[str]            =						processor(__snake_case				,		__snake_case				,		cur_len=__snake_case   )
												return scores

							UpperCAmelCase_    :      Optional[int]            =						jax.jit(__snake_case   )
							UpperCAmelCase_    :      Optional[Any]            =						jax.jit(__snake_case   )

							UpperCAmelCase_    :      Dict            =						jitted_run_no_processor_list(__snake_case				,		__snake_case				,		__snake_case   )
							UpperCAmelCase_    :      int            =						jitted_run_processor_list(__snake_case				,		__snake_case				,		__snake_case   )

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

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









import argparse
from argparse import Namespace

import torch
from torch import nn

from transformers import XGLMConfig, XGLMForCausalLM






def        UpperCamelCase_(							_snake_case   :  Union[str, Any]	):




			"""simple docstring"""






			__a        =[
			    'decoder.version',
			    'decoder.output_projection.weight',
			    '_float_tensor',
			    'decoder.embed_positions._float_tensor',
			]
			for k in ignore_keys:
						state_dict.pop(_snake_case       ,     _snake_case	)






def        UpperCamelCase_(							_snake_case   :  Any	):




			"""simple docstring"""






			__a      ,    __a        =emb.weight.shape
			__a        =nn.Linear(_snake_case       ,     _snake_case       ,     bias=_snake_case	)
			__a        =emb.weight.data
			return lin_layer






def        UpperCamelCase_(							_snake_case   :  str	):




			"""simple docstring"""






			__a        =torch.load(_snake_case       ,     map_location='cpu'	)
			__a        =Namespace(**checkpoint['cfg']['model']	)
			__a        =checkpoint['model']
			remove_ignore_keys_(_snake_case	)
			__a        =state_dict['decoder.embed_tokens.weight'].shape[0]

			__a        ={key.replace('decoder'       ,     'model'	): val for key, val in state_dict.items()}

			__a        =XGLMConfig(
			    vocab_size=_snake_case       ,     max_position_embeddings=args.max_target_positions       ,     num_layers=args.decoder_layers       ,     attention_heads=args.decoder_attention_heads       ,     ffn_dim=args.decoder_ffn_embed_dim       ,     d_model=args.decoder_embed_dim       ,     layerdrop=args.decoder_layerdrop       ,     dropout=args.dropout       ,     attention_dropout=args.attention_dropout       ,     activation_dropout=args.activation_dropout       ,     activation_function='gelu'       ,     scale_embedding=not args.no_scale_embedding       ,     tie_word_embeddings=args.share_decoder_input_output_embed       ,     )

			__a        =XGLMForCausalLM(_snake_case	)
			__a        =model.load_state_dict(_snake_case       ,     strict=_snake_case	)
			print(_snake_case	)
			__a        =make_linear_from_emb(model.model.embed_tokens	)

			return model


if __name__ == "__main__":
						_lowerCAmelCase					:			str  				=			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument("fairseq_path", type=str, help="path to a model.pt on local filesystem.")
						parser.add_argument("pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
						_lowerCAmelCase					:			str  				=			parser.parse_args()
						_lowerCAmelCase					:			Union[str, Any]  				=			convert_fairseq_xglm_checkpoint_from_disk(args.fairseq_path)
						model.save_pretrained(args.pytorch_dump_folder_path)






