






import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


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

							from transformers.modeling_flax_pytorch_utils import (
							    convert_pytorch_state_dict_to_flax,
							    load_flax_weights_in_pytorch_model,
							)
							from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
							import torch




@require_flax
class 						UpperCAmelCase_  :


						'''simple docstring'''




						def __init__(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE=14						, _SCREAMING_SNAKE_CASE=7						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE=99						, _SCREAMING_SNAKE_CASE=32						, _SCREAMING_SNAKE_CASE=5						, _SCREAMING_SNAKE_CASE=4						, _SCREAMING_SNAKE_CASE=37						, _SCREAMING_SNAKE_CASE="gelu"						, _SCREAMING_SNAKE_CASE=0.1						, _SCREAMING_SNAKE_CASE=0.1						, _SCREAMING_SNAKE_CASE=512						, _SCREAMING_SNAKE_CASE=0.02						, _SCREAMING_SNAKE_CASE=None						, )       -> Any:
								snake_case_ :							Dict         =				parent
								snake_case_ :							int         =				batch_size
								snake_case_ :							List[Any]         =				seq_length
								snake_case_ :							Dict         =				is_training
								snake_case_ :							str         =				use_input_mask
								snake_case_ :							str         =				use_labels
								snake_case_ :							Optional[int]         =				vocab_size
								snake_case_ :							List[Any]         =				d_model
								snake_case_ :							Optional[Any]         =				num_hidden_layers
								snake_case_ :							str         =				num_attention_heads
								snake_case_ :							List[str]         =				ffn_dim
								snake_case_ :							List[Any]         =				activation_function
								snake_case_ :							Any         =				activation_dropout
								snake_case_ :							List[str]         =				attention_dropout
								snake_case_ :							Any         =				max_position_embeddings
								snake_case_ :							Optional[Any]         =				initializer_range
								snake_case_ :							Union[str, Any]         =				None
								snake_case_ :							List[str]         =				0
								snake_case_ :							int         =				2
								snake_case_ :							List[str]         =				1






						def 					_lowerCAmelCase				(		self )       -> int:
								snake_case_ :							List[str]         =				np.clip(ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size )						, 3						, self.vocab_size )

								snake_case_ :							Optional[int]         =				None
								if self.use_input_mask:
										snake_case_ :							Any         =				random_attention_mask([self.batch_size, self.seq_length] )

								snake_case_ :							List[Any]         =				XGLMConfig(
								    vocab_size=self.vocab_size						, d_model=self.hidden_size						, num_layers=self.num_hidden_layers						, attention_heads=self.num_attention_heads						, ffn_dim=self.ffn_dim						, activation_function=self.activation_function						, activation_dropout=self.activation_dropout						, attention_dropout=self.attention_dropout						, max_position_embeddings=self.max_position_embeddings						, initializer_range=self.initializer_range						, use_cache=_SCREAMING_SNAKE_CASE						, bos_token_id=self.bos_token_id						, eos_token_id=self.eos_token_id						, pad_token_id=self.pad_token_id						, )

								return (config, input_ids, input_mask)






						def 					_lowerCAmelCase				(		self )       -> str:
								snake_case_ :							str         =				self.prepare_config_and_inputs()
								snake_case_       ,				snake_case_       ,				snake_case_ :							Any         =				config_and_inputs
								snake_case_ :							Dict         =				{"input_ids": input_ids, "attention_mask": attention_mask}
								return config, inputs_dict






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_       ,				snake_case_       ,				snake_case_ :							Tuple         =				self.prepare_config_and_inputs()

								snake_case_ :							Tuple         =				floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
								snake_case_ :							Optional[Any]         =				ids_tensor([self.batch_size, self.seq_length]						, vocab_size=2 )

								return (
								    config,
								    input_ids,
								    attention_mask,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> int:
								snake_case_ :							Optional[Any]         =				20
								snake_case_ :							List[str]         =				model_class_name(_SCREAMING_SNAKE_CASE )

								snake_case_ :							int         =				model.init_cache(input_ids.shape[0]						, _SCREAMING_SNAKE_CASE )
								snake_case_ :							Tuple         =				jnp.ones((input_ids.shape[0], max_decoder_length)						, dtype="i4" )

								snake_case_ :							Tuple         =				jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1 )[None, :]						, (input_ids.shape[0], input_ids.shape[-1] - 1) )
								snake_case_ :							Optional[int]         =				model(
								    input_ids[:, :-1]						, attention_mask=_SCREAMING_SNAKE_CASE						, past_key_values=_SCREAMING_SNAKE_CASE						, position_ids=_SCREAMING_SNAKE_CASE						, )

								snake_case_ :							Tuple         =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						, dtype="i4" )
								snake_case_ :							Union[str, Any]         =				model(
								    input_ids[:, -1:]						, attention_mask=_SCREAMING_SNAKE_CASE						, past_key_values=outputs_cache.past_key_values						, position_ids=_SCREAMING_SNAKE_CASE						, )

								snake_case_ :							Dict         =				model(_SCREAMING_SNAKE_CASE )

								snake_case_ :							List[Any]         =				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
								self.parent.assertTrue(diff < 1e-3						, msg=f'''Max diff is {diff}''' )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> int:
								snake_case_ :							Optional[Any]         =				20
								snake_case_ :							List[str]         =				model_class_name(_SCREAMING_SNAKE_CASE )

								snake_case_ :							Union[str, Any]         =				jnp.concatenate(
								    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1]) )]						, axis=-1						, )

								snake_case_ :							str         =				model.init_cache(input_ids.shape[0]						, _SCREAMING_SNAKE_CASE )
								snake_case_ :							Any         =				jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1 )[None, :]						, (input_ids.shape[0], input_ids.shape[-1] - 1) )

								snake_case_ :							Union[str, Any]         =				model(
								    input_ids[:, :-1]						, attention_mask=_SCREAMING_SNAKE_CASE						, past_key_values=_SCREAMING_SNAKE_CASE						, position_ids=_SCREAMING_SNAKE_CASE						, )
								snake_case_ :							Optional[int]         =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						, dtype="i4" )
								snake_case_ :							str         =				model(
								    input_ids[:, -1:]						, past_key_values=outputs_cache.past_key_values						, attention_mask=_SCREAMING_SNAKE_CASE						, position_ids=_SCREAMING_SNAKE_CASE						, )

								snake_case_ :							Any         =				model(_SCREAMING_SNAKE_CASE						, attention_mask=_SCREAMING_SNAKE_CASE )
								snake_case_ :							List[str]         =				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
								self.parent.assertTrue(diff < 1e-3						, msg=f'''Max diff is {diff}''' )




@require_sentencepiece
@require_flax
class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ,     SCREAMING_SNAKE_CASE__    ,     unittest.TestCase    ):


						'''simple docstring'''




						A   :      List[Any]						 =    (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
						A   :      Tuple						 =    (FlaxXGLMForCausalLM,) if is_flax_available() else ()






						def 					_lowerCAmelCase				(		self )       -> List[str]:
								snake_case_ :							Any         =				FlaxXGLMModelTester(self )






						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								for model_class_name in self.all_model_classes:
										snake_case_       ,				snake_case_       ,				snake_case_ :							Dict         =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> List[str]:
								for model_class_name in self.all_model_classes:
										snake_case_       ,				snake_case_       ,				snake_case_ :							Any         =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward_with_attn_mask(
										    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						@slow
						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							Union[str, Any]         =				XGLMTokenizer.from_pretrained("XGLM"						, padding_side="left" )
								snake_case_ :							List[str]         =				tokenizer(["Hello this is a long string", "Hey"]						, return_tensors="np"						, padding=_SCREAMING_SNAKE_CASE						, truncation=_SCREAMING_SNAKE_CASE )

								snake_case_ :							List[Any]         =				FlaxXGLMForCausalLM.from_pretrained("facebook/xglm-564M" )
								snake_case_ :							Optional[Any]         =				1
								snake_case_ :							int         =				False

								snake_case_ :							Optional[int]         =				jax.jit(model.generate )

								snake_case_ :							Optional[int]         =				jit_generate(inputs["input_ids"]						, attention_mask=inputs["attention_mask"] ).sequences

								snake_case_ :							Any         =				tokenizer.batch_decode(_SCREAMING_SNAKE_CASE						, skip_special_tokens=_SCREAMING_SNAKE_CASE )

								snake_case_ :							List[Any]         =				[
								    "Hello this is a long string of questions, but I'm not sure if I'm",
								    "Hey, I'm a newbie to the forum and I'",
								]

								self.assertListEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						@is_pt_flax_cross_test
						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_       ,				snake_case_ :							Dict         =				self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__ ):
												# prepare inputs
												snake_case_ :							Optional[int]         =				self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
												snake_case_ :							Dict         =				{k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

												# load corresponding PyTorch class
												snake_case_ :							Optional[int]         =				model_class.__name__[4:]  # Skip the "Flax" at the beginning
												snake_case_ :							List[Any]         =				getattr(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )

												snake_case_       ,				snake_case_ :							Any         =				pt_inputs["input_ids"].shape
												snake_case_ :							Optional[Any]         =				np.random.randint(0						, seq_length - 1						, size=(batch_size,) )
												for batch_idx, start_index in enumerate(_SCREAMING_SNAKE_CASE ):
														snake_case_ :							int         =				0
														snake_case_ :							Optional[Any]         =				1
														snake_case_ :							List[Any]         =				0
														snake_case_ :							str         =				1
												snake_case_ :							Dict         =				pt_model_class(_SCREAMING_SNAKE_CASE ).eval()
												# Flax models don't use the `use_cache` option and cache is not returned as a default.
												# So we disable `use_cache` here for PyTorch model.
												snake_case_ :							Optional[int]         =				False
												snake_case_ :							int         =				model_class(_SCREAMING_SNAKE_CASE						, dtype=jnp.floataa )
												snake_case_ :							Union[str, Any]         =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()						, _SCREAMING_SNAKE_CASE )
												snake_case_ :							List[Any]         =				fx_state

												with torch.no_grad():
														snake_case_ :							str         =				pt_model(**_SCREAMING_SNAKE_CASE ).to_tuple()

												snake_case_ :							List[str]         =				fx_model(**_SCREAMING_SNAKE_CASE ).to_tuple()
												self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, len(_SCREAMING_SNAKE_CASE )						, "Output lengths differ between Flax and PyTorch" )
												for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ):
														self.assert_almost_equals(fx_output[:, -1]						, pt_output[:, -1].numpy()						, 4e-2 )

												with tempfile.TemporaryDirectory() as tmpdirname:
														pt_model.save_pretrained(_SCREAMING_SNAKE_CASE )
														snake_case_ :							Any         =				model_class.from_pretrained(_SCREAMING_SNAKE_CASE						, from_pt=_SCREAMING_SNAKE_CASE )

												snake_case_ :							List[str]         =				fx_model_loaded(**_SCREAMING_SNAKE_CASE ).to_tuple()
												self.assertEqual(
												    len(_SCREAMING_SNAKE_CASE )						, len(_SCREAMING_SNAKE_CASE )						, "Output lengths differ between Flax and PyTorch" )
												for fx_output_loaded, pt_output in zip(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ):
														self.assert_almost_equals(fx_output_loaded[:, -1]						, pt_output[:, -1].numpy()						, 4e-2 )






						@is_pt_flax_cross_test
						def 					_lowerCAmelCase				(		self )       -> int:
								snake_case_       ,				snake_case_ :							Optional[Any]         =				self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__ ):
												# prepare inputs
												snake_case_ :							List[str]         =				self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
												snake_case_ :							Optional[Any]         =				{k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

												# load corresponding PyTorch class
												snake_case_ :							str         =				model_class.__name__[4:]  # Skip the "Flax" at the beginning
												snake_case_ :							int         =				getattr(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )

												snake_case_ :							Optional[int]         =				pt_model_class(_SCREAMING_SNAKE_CASE ).eval()
												snake_case_ :							Union[str, Any]         =				False
												snake_case_ :							Union[str, Any]         =				model_class(_SCREAMING_SNAKE_CASE						, dtype=jnp.floataa )

												snake_case_ :							Any         =				load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE						, fx_model.params )
												snake_case_       ,				snake_case_ :							Tuple         =				pt_inputs["input_ids"].shape
												snake_case_ :							Any         =				np.random.randint(0						, seq_length - 1						, size=(batch_size,) )
												for batch_idx, start_index in enumerate(_SCREAMING_SNAKE_CASE ):
														snake_case_ :							Optional[Any]         =				0
														snake_case_ :							Optional[int]         =				1
														snake_case_ :							Tuple         =				0
														snake_case_ :							Optional[int]         =				1

												# make sure weights are tied in PyTorch
												pt_model.tie_weights()

												with torch.no_grad():
														snake_case_ :							int         =				pt_model(**_SCREAMING_SNAKE_CASE ).to_tuple()

												snake_case_ :							int         =				fx_model(**_SCREAMING_SNAKE_CASE ).to_tuple()
												self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, len(_SCREAMING_SNAKE_CASE )						, "Output lengths differ between Flax and PyTorch" )
												for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ):
														self.assert_almost_equals(fx_output[:, -1]						, pt_output[:, -1].numpy()						, 4e-2 )

												with tempfile.TemporaryDirectory() as tmpdirname:
														fx_model.save_pretrained(_SCREAMING_SNAKE_CASE )
														snake_case_ :							int         =				pt_model_class.from_pretrained(_SCREAMING_SNAKE_CASE						, from_flax=_SCREAMING_SNAKE_CASE )

												with torch.no_grad():
														snake_case_ :							List[Any]         =				pt_model_loaded(**_SCREAMING_SNAKE_CASE ).to_tuple()

												self.assertEqual(
												    len(_SCREAMING_SNAKE_CASE )						, len(_SCREAMING_SNAKE_CASE )						, "Output lengths differ between Flax and PyTorch" )
												for fx_output, pt_output in zip(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ):
														self.assert_almost_equals(fx_output[:, -1]						, pt_output[:, -1].numpy()						, 4e-2 )






						@slow
						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								for model_class_name in self.all_model_classes:
										snake_case_ :							Optional[Any]         =				model_class_name.from_pretrained("facebook/xglm-564M" )
										snake_case_ :							Dict         =				model(np.ones((1, 1) ) )
										self.assertIsNotNone(_SCREAMING_SNAKE_CASE )








from __future__ import annotations

def    lowerCAmelCase__		(			_a				:	list[int]     ,					_a				:	int			):
		snake_case_       ,				snake_case_ :							Tuple         =				0, len(_a			) - 1
		while start_ind <= end_ind:
				if array[start_ind] == search_item:
						return start_ind
				elif array[end_ind] == search_item:
						return end_ind
				else:
						start_ind += 1
						end_ind -= 1
    # returns -1 if search_item is not found in array
		return -1


if __name__ == "__main__":
							print(double_linear_search(list(range(1_00)), 40))
