



import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available






def 			UpperCAmelCase__     (     lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__=None )			->   Tuple:
						__lowercase											=					None
						# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
						if str(weight.dtype ) == "torch.bfloat16":
												# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
												__lowercase											=					weight.view(torch.intaa )
												__lowercase											=					"""bfloat16"""
						__lowercase											=					weight.cpu().numpy()
						__lowercase											=					os.path.join(lowercase__    ,							F"{weight_name}.dat" )
						if index is not None:
												if dtype is None:
																		__lowercase											=					str(array.dtype )
												__lowercase											=					{"""dtype""": dtype, """shape""": list(array.shape )}
						if array.ndim == 0:
												__lowercase											=					array[None]
						__lowercase											=					np.memmap(lowercase__    ,							dtype=array.dtype    ,							mode="""w+"""    ,							shape=array.shape )
						__lowercase											=					array[:]
						file_array.flush()
						return index






def 			UpperCAmelCase__     (     lowercase__    ,							lowercase__ )			->   Dict:
						__lowercase											=					tuple(weight_info["""shape"""] )
						if shape == ():
												# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
												__lowercase											=					(1,)

						__lowercase											=					weight_info["""dtype"""]
						if dtype == "bfloat16":
												# NumPy does not support bfloat16 so this was saved as a int16
												__lowercase											=					"""int16"""

						__lowercase											=					np.memmap(lowercase__    ,							dtype=lowercase__    ,							shape=lowercase__    ,							mode="""r""" )

						if len(weight_info["""shape"""] ) == 0:
												__lowercase											=					weight[0]
						__lowercase											=					torch.tensor(lowercase__ )
						if weight_info["dtype"] == "bfloat16":
												__lowercase											=					weight.view(torch.bfloataa )

						return weight






def 			UpperCAmelCase__     (     lowercase__    ,							lowercase__ )			->   int:
						if index is None or len(lowercase__ ) == 0:
												# Nothing to save
												return

						__lowercase											=					os.path.join(lowercase__    ,							"""index.json""" )
						if os.path.isfile(lowercase__ ):
												with open(lowercase__    ,							"""r"""    ,							encoding="""utf-8""" ) as f:
																		__lowercase											=					json.load(lowercase__ )
						else:
												__lowercase											=					{}
						current_index.update(lowercase__ )

						with open(lowercase__    ,							"""w"""    ,							encoding="""utf-8""" ) as f:
												json.dump(lowercase__    ,							lowercase__    ,							indent=2 )






def 			UpperCAmelCase__     (     lowercase__    ,							lowercase__ )			->   Optional[int]:
						os.makedirs(lowercase__    ,							exist_ok=lowercase__ )
						__lowercase											=					{}
						for name, parameter in state_dict.items():
												__lowercase											=					offload_weight(lowercase__    ,							lowercase__    ,							lowercase__    ,							index=lowercase__ )

						# Update index
						save_offload_index(lowercase__    ,							lowercase__ )
class       _lowerCAmelCase (			_UpperCAmelCase      ):






			"""simple docstring"""





			def __init__(					self						:							Union[str, Any]      ,  lowercase						:							str      ,  lowercase						:							Union[str, Any] )    ->       int:
									"""simple docstring"""







									__lowercase											=					dataset
									__lowercase											=					prefix





			def __getitem__(					self						:							str      ,  lowercase						:							List[str] )    ->       Any:
									"""simple docstring"""







									return self.dataset[F"{self.prefix}{key}"]





			def __iter__(					self						:							int )    ->       Any:
									"""simple docstring"""







									return iter([key for key in self.dataset if key.startswith(self.prefix )] )







			def __len__(					self						:							Optional[int] )    ->       int:
									"""simple docstring"""







									return len(self.dataset )





class       _lowerCAmelCase (			_UpperCAmelCase      ):






			"""simple docstring"""





			def __init__(					self						:							Any      ,  lowercase						:							Any = None      ,  lowercase						:							Union[str, Any] = None      ,  lowercase						:							int = None      ,  lowercase						:							List[Any]=None      ,  )    ->       str:
									"""simple docstring"""







									if state_dict is None and save_folder is None:
															raise ValueError("""Need either a `state_dict` or a `save_folder` containing offloaded weights.""" )

									__lowercase											=					{} if state_dict is None else state_dict
									__lowercase											=					save_folder
									if index is None and save_folder is not None:
															with open(os.path.join(_UpperCAmelCase      ,  """index.json""" ) ) as f:
																					__lowercase											=					json.load(_UpperCAmelCase )
									__lowercase											=					{} if index is None else index
									__lowercase											=					list(self.state_dict.keys() )
									self.all_keys.extend([key for key in self.index if key not in self.all_keys] )
									__lowercase											=					device





			def __getitem__(					self						:							int      ,  lowercase						:							int )    ->       Tuple:
									"""simple docstring"""







									if key in self.state_dict:
															return self.state_dict[key]
									__lowercase											=					self.index[key]
									if weight_info.get("""safetensors_file""" ) is not None:
															if not is_safetensors_available():
																					raise ImportError("""These offloaded weights require the use of safetensors: `pip install safetensors`.""" )

															from safetensors import safe_open

															__lowercase											=					"""cpu""" if self.device is None else self.device
															with safe_open(weight_info["""safetensors_file"""]      ,  framework="""pt"""      ,  device=_UpperCAmelCase ) as f:
																					__lowercase											=					f.get_tensor(weight_info.get("""weight_name"""      ,  _UpperCAmelCase ) )

															if "dtype" in weight_info:
																					return tensor.to(getattr(_UpperCAmelCase      ,  weight_info["""dtype"""] ) )
															else:
																					return tensor

									__lowercase											=					os.path.join(self.save_folder      ,  F"{key}.dat" )
									return load_offloaded_weight(_UpperCAmelCase      ,  _UpperCAmelCase )





			def __iter__(					self						:							Optional[int] )    ->       Optional[int]:
									"""simple docstring"""







									return iter(self.all_keys )







			def __len__(					self						:							Dict )    ->       Any:
									"""simple docstring"""







									return len(self.all_keys )






def 			UpperCAmelCase__     (     lowercase__    ,							lowercase__ )			->   Dict:
						__lowercase											=					{}
						for module_name in submodule_names:
												# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
												# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
												result.update(
												    {
												        key: param
												        for key, param in state_dict.items()
												        if key == module_name or key.startswith(module_name + """.""" )
												    } )
						return result













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 		lowerCAmelCase_				:





     def __init__(	self							, _UpperCAmelCase							, _UpperCAmelCase=14							, _UpperCAmelCase=7							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase=99							, _UpperCAmelCase=32							, _UpperCAmelCase=5							, _UpperCAmelCase=4							, _UpperCAmelCase=37							, _UpperCAmelCase="gelu"							, _UpperCAmelCase=0.1							, _UpperCAmelCase=0.1							, _UpperCAmelCase=512							, _UpperCAmelCase=0.02							, _UpperCAmelCase=None							, )		->				str:
        lowerCamelCase									=parent
        lowerCamelCase									=batch_size
        lowerCamelCase									=seq_length
        lowerCamelCase									=is_training
        lowerCamelCase									=use_input_mask
        lowerCamelCase									=use_labels
        lowerCamelCase									=vocab_size
        lowerCamelCase									=d_model
        lowerCamelCase									=num_hidden_layers
        lowerCamelCase									=num_attention_heads
        lowerCamelCase									=ffn_dim
        lowerCamelCase									=activation_function
        lowerCamelCase									=activation_dropout
        lowerCamelCase									=attention_dropout
        lowerCamelCase									=max_position_embeddings
        lowerCamelCase									=initializer_range
        lowerCamelCase									=None
        lowerCamelCase									=0
        lowerCamelCase									=2
        lowerCamelCase									=1





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

        lowerCamelCase									=None
        if self.use_input_mask:
           lowerCamelCase									=random_attention_mask([self.batch_size, self.seq_length]  )

        lowerCamelCase									=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=_UpperCAmelCase							, 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  )		->				Tuple:
        lowerCamelCase									=self.prepare_config_and_inputs()
        lowerCamelCase				,       lowerCamelCase				,       lowerCamelCase									=config_and_inputs
        lowerCamelCase									={"""input_ids""": input_ids, """attention_mask""": attention_mask}
        return config, inputs_dict





     def 						lowerCamelCase							(	self  )		->				int:
        lowerCamelCase				,       lowerCamelCase				,       lowerCamelCase									=self.prepare_config_and_inputs()

        lowerCamelCase									=floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
        lowerCamelCase									=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							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )		->				Optional[Any]:
        lowerCamelCase									=20
        lowerCamelCase									=model_class_name(_UpperCAmelCase  )

        lowerCamelCase									=model.init_cache(input_ids.shape[0]							, _UpperCAmelCase  )
        lowerCamelCase									=jnp.ones((input_ids.shape[0], max_decoder_length)							, dtype="""i4"""  )

        lowerCamelCase									=jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1  )[None, :]							, (input_ids.shape[0], input_ids.shape[-1] - 1)  )
        lowerCamelCase									=model(
            input_ids[:, :-1]							, attention_mask=_UpperCAmelCase							, past_key_values=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )

        lowerCamelCase									=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							, dtype="""i4"""  )
        lowerCamelCase									=model(
            input_ids[:, -1:]							, attention_mask=_UpperCAmelCase							, past_key_values=outputs_cache.past_key_values							, position_ids=_UpperCAmelCase							, )

        lowerCamelCase									=model(_UpperCAmelCase  )

        lowerCamelCase									=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							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )		->				List[Any]:
        lowerCamelCase									=20
        lowerCamelCase									=model_class_name(_UpperCAmelCase  )

        lowerCamelCase									=jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])  )]							, axis=-1							, )

        lowerCamelCase									=model.init_cache(input_ids.shape[0]							, _UpperCAmelCase  )
        lowerCamelCase									=jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1  )[None, :]							, (input_ids.shape[0], input_ids.shape[-1] - 1)  )

        lowerCamelCase									=model(
            input_ids[:, :-1]							, attention_mask=_UpperCAmelCase							, past_key_values=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )
        lowerCamelCase									=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							, dtype="""i4"""  )
        lowerCamelCase									=model(
            input_ids[:, -1:]							, past_key_values=outputs_cache.past_key_values							, attention_mask=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )

        lowerCamelCase									=model(_UpperCAmelCase							, attention_mask=_UpperCAmelCase  )
        lowerCamelCase									=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 		lowerCAmelCase_				(  UpperCamelCase				,     UpperCamelCase				,     unittest.TestCase ):
     __lowerCamelCase       :    Tuple 	=     (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
     __lowerCamelCase       :    Dict 	=     (FlaxXGLMForCausalLM,) if is_flax_available() else ()





     def 						lowerCamelCase							(	self  )		->				List[str]:
        lowerCamelCase									=FlaxXGLMModelTester(self  )





     def 						lowerCamelCase							(	self  )		->				List[str]:
        for model_class_name in self.all_model_classes:
           lowerCamelCase				,       lowerCamelCase				,       lowerCamelCase									=self.model_tester.prepare_config_and_inputs()
           self.model_tester.check_use_cache_forward(_UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )





     def 						lowerCamelCase							(	self  )		->				Tuple:
        for model_class_name in self.all_model_classes:
           lowerCamelCase				,       lowerCamelCase				,       lowerCamelCase									=self.model_tester.prepare_config_and_inputs()
           self.model_tester.check_use_cache_forward_with_attn_mask(
               _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )





     @slow
     def 						lowerCamelCase							(	self  )		->				Tuple:
        lowerCamelCase									=XGLMTokenizer.from_pretrained("""XGLM"""							, padding_side="""left"""  )
        lowerCamelCase									=tokenizer(["""Hello this is a long string""", """Hey"""]							, return_tensors="""np"""							, padding=_UpperCAmelCase							, truncation=_UpperCAmelCase  )

        lowerCamelCase									=FlaxXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""  )
        lowerCamelCase									=1
        lowerCamelCase									=False

        lowerCamelCase									=jax.jit(model.generate  )

        lowerCamelCase									=jit_generate(inputs["""input_ids"""]							, attention_mask=inputs["""attention_mask"""]  ).sequences

        lowerCamelCase									=tokenizer.batch_decode(_UpperCAmelCase							, skip_special_tokens=_UpperCAmelCase  )

        lowerCamelCase									=[
            """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(_UpperCAmelCase							, _UpperCAmelCase  )





     @is_pt_flax_cross_test
     def 						lowerCamelCase							(	self  )		->				Dict:
        lowerCamelCase				,       lowerCamelCase									=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
              lowerCamelCase									=self._prepare_for_class(_UpperCAmelCase							, _UpperCAmelCase  )
              lowerCamelCase									={k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

              # load corresponding PyTorch class
              lowerCamelCase									=model_class.__name__[4:]  # Skip the "Flax" at the beginning
              lowerCamelCase									=getattr(_UpperCAmelCase							, _UpperCAmelCase  )

              lowerCamelCase				,       lowerCamelCase									=pt_inputs["""input_ids"""].shape
              lowerCamelCase									=np.random.randint(0							, seq_length - 1							, size=(batch_size,)  )
              for batch_idx, start_index in enumerate(_UpperCAmelCase  ):
                 lowerCamelCase									=0
                 lowerCamelCase									=1
                 lowerCamelCase									=0
                 lowerCamelCase									=1
              lowerCamelCase									=pt_model_class(_UpperCAmelCase  ).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.
              lowerCamelCase									=False
              lowerCamelCase									=model_class(_UpperCAmelCase							, dtype=jnp.floataa  )
              lowerCamelCase									=convert_pytorch_state_dict_to_flax(pt_model.state_dict()							, _UpperCAmelCase  )
              lowerCamelCase									=fx_state

              with torch.no_grad():
                 lowerCamelCase									=pt_model(**_UpperCAmelCase  ).to_tuple()

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

              with tempfile.TemporaryDirectory() as tmpdirname:
                 pt_model.save_pretrained(_UpperCAmelCase  )
                 lowerCamelCase									=model_class.from_pretrained(_UpperCAmelCase							, from_pt=_UpperCAmelCase  )

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





     @is_pt_flax_cross_test
     def 						lowerCamelCase							(	self  )		->				Optional[Any]:
        lowerCamelCase				,       lowerCamelCase									=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
              lowerCamelCase									=self._prepare_for_class(_UpperCAmelCase							, _UpperCAmelCase  )
              lowerCamelCase									={k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

              # load corresponding PyTorch class
              lowerCamelCase									=model_class.__name__[4:]  # Skip the "Flax" at the beginning
              lowerCamelCase									=getattr(_UpperCAmelCase							, _UpperCAmelCase  )

              lowerCamelCase									=pt_model_class(_UpperCAmelCase  ).eval()
              lowerCamelCase									=False
              lowerCamelCase									=model_class(_UpperCAmelCase							, dtype=jnp.floataa  )

              lowerCamelCase									=load_flax_weights_in_pytorch_model(_UpperCAmelCase							, fx_model.params  )
              lowerCamelCase				,       lowerCamelCase									=pt_inputs["""input_ids"""].shape
              lowerCamelCase									=np.random.randint(0							, seq_length - 1							, size=(batch_size,)  )
              for batch_idx, start_index in enumerate(_UpperCAmelCase  ):
                 lowerCamelCase									=0
                 lowerCamelCase									=1
                 lowerCamelCase									=0
                 lowerCamelCase									=1

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

              with torch.no_grad():
                 lowerCamelCase									=pt_model(**_UpperCAmelCase  ).to_tuple()

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

              with tempfile.TemporaryDirectory() as tmpdirname:
                 fx_model.save_pretrained(_UpperCAmelCase  )
                 lowerCamelCase									=pt_model_class.from_pretrained(_UpperCAmelCase							, from_flax=_UpperCAmelCase  )

              with torch.no_grad():
                 lowerCamelCase									=pt_model_loaded(**_UpperCAmelCase  ).to_tuple()

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





     @slow
     def 						lowerCamelCase							(	self  )		->				int:
        for model_class_name in self.all_model_classes:
           lowerCamelCase									=model_class_name.from_pretrained("""facebook/xglm-564M"""  )
           lowerCamelCase									=model(np.ones((1, 1)  )  )
           self.assertIsNotNone(_UpperCAmelCase  )