code: 
  
  
  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  :				Dict		,				SCREAMING_SNAKE_CASE_  :				Any		,				SCREAMING_SNAKE_CASE_  :				List[Any]=1_4		,				SCREAMING_SNAKE_CASE_  :				Dict=7		,				SCREAMING_SNAKE_CASE_  :				List[Any]=True		,				SCREAMING_SNAKE_CASE_  :				List[str]=True		,				SCREAMING_SNAKE_CASE_  :				List[Any]=True		,				SCREAMING_SNAKE_CASE_  :				int=9_9		,				SCREAMING_SNAKE_CASE_  :				Dict=3_2		,				SCREAMING_SNAKE_CASE_  :				str=5		,				SCREAMING_SNAKE_CASE_  :				List[Any]=4		,				SCREAMING_SNAKE_CASE_  :				List[Any]=3_7		,				SCREAMING_SNAKE_CASE_  :				List[Any]="gelu"		,				SCREAMING_SNAKE_CASE_  :				str=0.1		,				SCREAMING_SNAKE_CASE_  :				List[str]=0.1		,				SCREAMING_SNAKE_CASE_  :				Dict=5_1_2		,				SCREAMING_SNAKE_CASE_  :				Optional[Any]=0.02		,				SCREAMING_SNAKE_CASE_  :				int=None		,				):
  													_a     =			parent
  													_a     =			batch_size
  													_a     =			seq_length
  													_a     =			is_training
  													_a     =			use_input_mask
  													_a     =			use_labels
  													_a     =			vocab_size
  													_a     =			d_model
  													_a     =			num_hidden_layers
  													_a     =			num_attention_heads
  													_a     =			ffn_dim
  													_a     =			activation_function
  													_a     =			activation_dropout
  													_a     =			attention_dropout
  													_a     =			max_position_embeddings
  													_a     =			initializer_range
  													_a     =			None
  													_a     =			0
  													_a     =			2
  													_a     =			1
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				int		):
  													_a     =			np.clip(ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size		)		,				3		,				self.vocab_size		)
  
  													_a     =			None
  													if self.use_input_mask:
  																			_a     =			random_attention_mask([self.batch_size, self.seq_length]		)
  
  													_a     =			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 		_UpperCAmelCase	(		self  :				Optional[Any]		):
  													_a     =			self.prepare_config_and_inputs()
  													_a ,	_a ,	_a     =			config_and_inputs
  													_a     =			{'input_ids': input_ids, 'attention_mask': attention_mask}
  													return config, inputs_dict
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				str		):
  													_a ,	_a ,	_a     =			self.prepare_config_and_inputs()
  
  													_a     =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
  													_a     =			ids_tensor([self.batch_size, self.seq_length]		,				vocab_size=2		)
  
  													return (
  													    config,
  													    input_ids,
  													    attention_mask,
  													    encoder_hidden_states,
  													    encoder_attention_mask,
  													)
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				List[Any]		,				SCREAMING_SNAKE_CASE_  :				Union[str, Any]		,				SCREAMING_SNAKE_CASE_  :				Any		,				SCREAMING_SNAKE_CASE_  :				Dict		,				SCREAMING_SNAKE_CASE_  :				Optional[int]		):
  													_a     =			2_0
  													_a     =			model_class_name(SCREAMING_SNAKE_CASE_		)
  
  													_a     =			model.init_cache(input_ids.shape[0]		,				SCREAMING_SNAKE_CASE_		)
  													_a     =			jnp.ones((input_ids.shape[0], max_decoder_length)		,				dtype='i4'		)
  
  													_a     =			jnp.broadcast_to(
  													    jnp.arange(input_ids.shape[-1] - 1		)[None, :]		,				(input_ids.shape[0], input_ids.shape[-1] - 1)		)
  													_a     =			model(
  													    input_ids[:, :-1]		,				attention_mask=SCREAMING_SNAKE_CASE_		,				past_key_values=SCREAMING_SNAKE_CASE_		,				position_ids=SCREAMING_SNAKE_CASE_		,				)
  
  													_a     =			jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]		,				dtype='i4'		)
  													_a     =			model(
  													    input_ids[:, -1:]		,				attention_mask=SCREAMING_SNAKE_CASE_		,				past_key_values=outputs_cache.past_key_values		,				position_ids=SCREAMING_SNAKE_CASE_		,				)
  
  													_a     =			model(SCREAMING_SNAKE_CASE_		)
  
  													_a     =			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 		_UpperCAmelCase	(		self  :				int		,				SCREAMING_SNAKE_CASE_  :				Tuple		,				SCREAMING_SNAKE_CASE_  :				str		,				SCREAMING_SNAKE_CASE_  :				Optional[Any]		,				SCREAMING_SNAKE_CASE_  :				Any		):
  													_a     =			2_0
  													_a     =			model_class_name(SCREAMING_SNAKE_CASE_		)
  
  													_a     =			jnp.concatenate(
  													    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])		)]		,				axis=-1		,				)
  
  													_a     =			model.init_cache(input_ids.shape[0]		,				SCREAMING_SNAKE_CASE_		)
  													_a     =			jnp.broadcast_to(
  													    jnp.arange(input_ids.shape[-1] - 1		)[None, :]		,				(input_ids.shape[0], input_ids.shape[-1] - 1)		)
  
  													_a     =			model(
  													    input_ids[:, :-1]		,				attention_mask=SCREAMING_SNAKE_CASE_		,				past_key_values=SCREAMING_SNAKE_CASE_		,				position_ids=SCREAMING_SNAKE_CASE_		,				)
  													_a     =			jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]		,				dtype='i4'		)
  													_a     =			model(
  													    input_ids[:, -1:]		,				past_key_values=outputs_cache.past_key_values		,				attention_mask=SCREAMING_SNAKE_CASE_		,				position_ids=SCREAMING_SNAKE_CASE_		,				)
  
  													_a     =			model(SCREAMING_SNAKE_CASE_		,				attention_mask=SCREAMING_SNAKE_CASE_		)
  													_a     =			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 (      lowerCamelCase__			,							lowerCamelCase__			,							unittest.TestCase	):
  
  
  
  
  
  
  
  							'''simple docstring'''
  							_A  				=  (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
  							_A  				=  (FlaxXGLMForCausalLM,) if is_flax_available() else ()
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				Optional[int]		):
  													_a     =			FlaxXGLMModelTester(self		)
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				List[Any]		):
  													for model_class_name in self.all_model_classes:
  																			_a ,	_a ,	_a     =			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 		_UpperCAmelCase	(		self  :				Any		):
  													for model_class_name in self.all_model_classes:
  																			_a ,	_a ,	_a     =			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 		_UpperCAmelCase	(		self  :				Union[str, Any]		):
  													_a     =			XGLMTokenizer.from_pretrained('XGLM'		,				padding_side='left'		)
  													_a     =			tokenizer(['Hello this is a long string', 'Hey']		,				return_tensors='np'		,				padding=SCREAMING_SNAKE_CASE_		,				truncation=SCREAMING_SNAKE_CASE_		)
  
  													_a     =			FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M'		)
  													_a     =			1
  													_a     =			False
  
  													_a     =			jax.jit(model.generate		)
  
  													_a     =			jit_generate(inputs['input_ids']		,				attention_mask=inputs['attention_mask']		).sequences
  
  													_a     =			tokenizer.batch_decode(SCREAMING_SNAKE_CASE_		,				skip_special_tokens=SCREAMING_SNAKE_CASE_		)
  
  													_a     =			[
  													    '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 		_UpperCAmelCase	(		self  :				int		):
  													_a ,	_a     =			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
  																									_a     =			self._prepare_for_class(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  																									_a     =			{k: torch.tensor(v.tolist()		) for k, v in prepared_inputs_dict.items()}
  
  																									# load corresponding PyTorch class
  																									_a     =			model_class.__name__[4:]  # Skip the "Flax" at the beginning
  																									_a     =			getattr(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  
  																									_a ,	_a     =			pt_inputs['input_ids'].shape
  																									_a     =			np.random.randint(0		,				seq_length - 1		,				size=(batch_size,)		)
  																									for batch_idx, start_index in enumerate(SCREAMING_SNAKE_CASE_		):
  																															_a     =			0
  																															_a     =			1
  																															_a     =			0
  																															_a     =			1
  																									_a     =			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.
  																									_a     =			False
  																									_a     =			model_class(SCREAMING_SNAKE_CASE_		,				dtype=jnp.floataa		)
  																									_a     =			convert_pytorch_state_dict_to_flax(pt_model.state_dict()		,				SCREAMING_SNAKE_CASE_		)
  																									_a     =			fx_state
  
  																									with torch.no_grad():
  																															_a     =			pt_model(**SCREAMING_SNAKE_CASE_		).to_tuple()
  
  																									_a     =			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_		)
  																															_a     =			model_class.from_pretrained(SCREAMING_SNAKE_CASE_		,				from_pt=SCREAMING_SNAKE_CASE_		)
  
  																									_a     =			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 		_UpperCAmelCase	(		self  :				List[str]		):
  													_a ,	_a     =			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
  																									_a     =			self._prepare_for_class(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  																									_a     =			{k: torch.tensor(v.tolist()		) for k, v in prepared_inputs_dict.items()}
  
  																									# load corresponding PyTorch class
  																									_a     =			model_class.__name__[4:]  # Skip the "Flax" at the beginning
  																									_a     =			getattr(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  
  																									_a     =			pt_model_class(SCREAMING_SNAKE_CASE_		).eval()
  																									_a     =			False
  																									_a     =			model_class(SCREAMING_SNAKE_CASE_		,				dtype=jnp.floataa		)
  
  																									_a     =			load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE_		,				fx_model.params		)
  																									_a ,	_a     =			pt_inputs['input_ids'].shape
  																									_a     =			np.random.randint(0		,				seq_length - 1		,				size=(batch_size,)		)
  																									for batch_idx, start_index in enumerate(SCREAMING_SNAKE_CASE_		):
  																															_a     =			0
  																															_a     =			1
  																															_a     =			0
  																															_a     =			1
  
  																									# make sure weights are tied in PyTorch
  																									pt_model.tie_weights()
  
  																									with torch.no_grad():
  																															_a     =			pt_model(**SCREAMING_SNAKE_CASE_		).to_tuple()
  
  																									_a     =			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_		)
  																															_a     =			pt_model_class.from_pretrained(SCREAMING_SNAKE_CASE_		,				from_flax=SCREAMING_SNAKE_CASE_		)
  
  																									with torch.no_grad():
  																															_a     =			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 		_UpperCAmelCase	(		self  :				int		):
  													for model_class_name in self.all_model_classes:
  																			_a     =			model_class_name.from_pretrained('facebook/xglm-564M'		)
  																			_a     =			model(np.ones((1, 1)		)		)
  																			self.assertIsNotNone(SCREAMING_SNAKE_CASE_		)
  
  
  
style_context: 
  
  
  from typing import List, Union
  
  from ...configuration_utils import PretrainedConfig
  from ...utils import logging
  
  
  lowercase_								=    logging.get_logger(__name__)
  
  lowercase_								=    {
      'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/config.json',
      # See all LED models at https://huggingface.co/models?filter=led
  }
  
  
  
  
  class 					_UpperCamelCase (      lowerCamelCase__	):
  
  
  
  
  
  
  
  							'''simple docstring'''
  							_A  				=  "led"
  							_A  				=  {
  							    "num_attention_heads": "encoder_attention_heads",
  							    "hidden_size": "d_model",
  							    "attention_probs_dropout_prob": "attention_dropout",
  							    "initializer_range": "init_std",
  							}
  
  
  
  
  
  
  
  							def __init__(		self  :				List[str]		,				SCREAMING_SNAKE_CASE_  :				List[str]=5_0_2_6_5		,				SCREAMING_SNAKE_CASE_  :				str=1_6_3_8_4		,				SCREAMING_SNAKE_CASE_  :				List[Any]=1_0_2_4		,				SCREAMING_SNAKE_CASE_  :				List[str]=1_2		,				SCREAMING_SNAKE_CASE_  :				Any=4_0_9_6		,				SCREAMING_SNAKE_CASE_  :				Optional[Any]=1_6		,				SCREAMING_SNAKE_CASE_  :				str=1_2		,				SCREAMING_SNAKE_CASE_  :				List[Any]=4_0_9_6		,				SCREAMING_SNAKE_CASE_  :				List[str]=1_6		,				SCREAMING_SNAKE_CASE_  :				Optional[int]=0.0		,				SCREAMING_SNAKE_CASE_  :				Union[str, Any]=0.0		,				SCREAMING_SNAKE_CASE_  :				List[Any]=True		,				SCREAMING_SNAKE_CASE_  :				str=True		,				SCREAMING_SNAKE_CASE_  :				List[str]="gelu"		,				SCREAMING_SNAKE_CASE_  :				List[Any]=1_0_2_4		,				SCREAMING_SNAKE_CASE_  :				Tuple=0.1		,				SCREAMING_SNAKE_CASE_  :				Dict=0.0		,				SCREAMING_SNAKE_CASE_  :				Optional[Any]=0.0		,				SCREAMING_SNAKE_CASE_  :				Tuple=0.02		,				SCREAMING_SNAKE_CASE_  :				str=2		,				SCREAMING_SNAKE_CASE_  :				Tuple=0.0		,				SCREAMING_SNAKE_CASE_  :				Optional[int]=1		,				SCREAMING_SNAKE_CASE_  :				str=0		,				SCREAMING_SNAKE_CASE_  :				Union[str, Any]=2		,				SCREAMING_SNAKE_CASE_  :				Union[List[int], int] = 5_1_2		,				**SCREAMING_SNAKE_CASE_  :				Dict		,				):
  													_a     =			vocab_size
  													_a     =			max_encoder_position_embeddings
  													_a     =			max_decoder_position_embeddings
  													_a     =			d_model
  													_a     =			encoder_ffn_dim
  													_a     =			encoder_layers
  													_a     =			encoder_attention_heads
  													_a     =			decoder_ffn_dim
  													_a     =			decoder_layers
  													_a     =			decoder_attention_heads
  													_a     =			dropout
  													_a     =			attention_dropout
  													_a     =			activation_dropout
  													_a     =			activation_function
  													_a     =			init_std
  													_a     =			encoder_layerdrop
  													_a     =			decoder_layerdrop
  													_a     =			classifier_dropout
  													_a     =			use_cache
  													_a     =			encoder_layers
  													_a     =			attention_window
  
  													super().__init__(
  													    pad_token_id=SCREAMING_SNAKE_CASE_		,				bos_token_id=SCREAMING_SNAKE_CASE_		,				eos_token_id=SCREAMING_SNAKE_CASE_		,				is_encoder_decoder=SCREAMING_SNAKE_CASE_		,				decoder_start_token_id=SCREAMING_SNAKE_CASE_		,				**SCREAMING_SNAKE_CASE_		,				)
  
  
  