
import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import (
				    RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
				    RwkvForCausalLM,
				    RwkvModel,
				)
				from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
				lowercase__							:							int													=     False






class     SCREAMING_SNAKE_CASE__      :

							"""simple docstring"""






							def __init__(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_=14       ,		SCREAMING_SNAKE_CASE_=7       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_=False       ,		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_=37       ,		SCREAMING_SNAKE_CASE_="gelu"       ,		SCREAMING_SNAKE_CASE_=0.1       ,		SCREAMING_SNAKE_CASE_=0.1       ,		SCREAMING_SNAKE_CASE_=512       ,		SCREAMING_SNAKE_CASE_=16       ,		SCREAMING_SNAKE_CASE_=2       ,		SCREAMING_SNAKE_CASE_=3       ,		SCREAMING_SNAKE_CASE_=4       ,		SCREAMING_SNAKE_CASE_=None       ,		)->				Tuple:





													'''simple docstring'''






													__UpperCamelCase       					=     parent
													__UpperCamelCase       					=     batch_size
													__UpperCamelCase       					=     seq_length
													__UpperCamelCase       					=     is_training
													__UpperCamelCase       					=     use_token_type_ids
													__UpperCamelCase       					=     use_input_mask
													__UpperCamelCase       					=     use_labels
													__UpperCamelCase       					=     use_mc_token_ids
													__UpperCamelCase       					=     vocab_size
													__UpperCamelCase       					=     hidden_size
													__UpperCamelCase       					=     num_hidden_layers
													__UpperCamelCase       					=     intermediate_size
													__UpperCamelCase       					=     hidden_act
													__UpperCamelCase       					=     hidden_dropout_prob
													__UpperCamelCase       					=     attention_probs_dropout_prob
													__UpperCamelCase       					=     max_position_embeddings
													__UpperCamelCase       					=     type_vocab_size
													__UpperCamelCase       					=     type_sequence_label_size
													__UpperCamelCase       					=     num_labels
													__UpperCamelCase       					=     num_choices
													__UpperCamelCase       					=     scope
													__UpperCamelCase       					=     vocab_size - 1
													__UpperCamelCase       					=     vocab_size - 1
													__UpperCamelCase       					=     vocab_size - 1



							def       A__						(					self     )->				str:





													'''simple docstring'''






													return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_=False       ,		SCREAMING_SNAKE_CASE_=False       ,		SCREAMING_SNAKE_CASE_=False     )->				str:





													'''simple docstring'''






													__UpperCamelCase       					=     ids_tensor([self.batch_size, self.seq_length]       ,		self.vocab_size     )

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

													__UpperCamelCase       					=     None
													if self.use_token_type_ids:
																			__UpperCamelCase       					=     ids_tensor([self.batch_size, self.seq_length]       ,		self.type_vocab_size     )

													__UpperCamelCase       					=     None
													if self.use_mc_token_ids:
																			__UpperCamelCase       					=     ids_tensor([self.batch_size, self.num_choices]       ,		self.seq_length     )

													__UpperCamelCase       					=     None
													__UpperCamelCase       					=     None
													__UpperCamelCase       					=     None
													if self.use_labels:
																			__UpperCamelCase       					=     ids_tensor([self.batch_size]       ,		self.type_sequence_label_size     )
																			__UpperCamelCase       					=     ids_tensor([self.batch_size, self.seq_length]       ,		self.num_labels     )
																			__UpperCamelCase       					=     ids_tensor([self.batch_size]       ,		self.num_choices     )

													__UpperCamelCase       					=     self.get_config(
													    gradient_checkpointing=SCREAMING_SNAKE_CASE_       ,		scale_attn_by_inverse_layer_idx=SCREAMING_SNAKE_CASE_       ,		reorder_and_upcast_attn=SCREAMING_SNAKE_CASE_       ,		)

													return (
													    config,
													    input_ids,
													    input_mask,
													    None,
													    token_type_ids,
													    mc_token_ids,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													)



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_=False       ,		SCREAMING_SNAKE_CASE_=False       ,		SCREAMING_SNAKE_CASE_=False     )->				Dict:





													'''simple docstring'''






													return RwkvConfig(
													    vocab_size=self.vocab_size       ,		hidden_size=self.hidden_size       ,		num_hidden_layers=self.num_hidden_layers       ,		intermediate_size=self.intermediate_size       ,		activation_function=self.hidden_act       ,		resid_pdrop=self.hidden_dropout_prob       ,		attn_pdrop=self.attention_probs_dropout_prob       ,		n_positions=self.max_position_embeddings       ,		type_vocab_size=self.type_vocab_size       ,		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       ,		gradient_checkpointing=SCREAMING_SNAKE_CASE_       ,		scale_attn_by_inverse_layer_idx=SCREAMING_SNAKE_CASE_       ,		reorder_and_upcast_attn=SCREAMING_SNAKE_CASE_       ,		)



							def       A__						(					self     )->				Optional[int]:





													'''simple docstring'''






													__UpperCamelCase       					=     self.get_config()
													__UpperCamelCase       					=     300
													return config



							def       A__						(					self     )->				Optional[int]:





													'''simple docstring'''






													(
													    (
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,														
													)       					=     self.prepare_config_and_inputs()

													__UpperCamelCase       					=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
													__UpperCamelCase       					=     ids_tensor([self.batch_size, self.seq_length]       ,		vocab_size=2     )

													return (
													    config,
													    input_ids,
													    input_mask,
													    head_mask,
													    token_type_ids,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													    encoder_hidden_states,
													    encoder_attention_mask,
													)



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		*SCREAMING_SNAKE_CASE_     )->				str:





													'''simple docstring'''






													__UpperCamelCase       					=     True
													__UpperCamelCase       					=     RwkvModel(config=SCREAMING_SNAKE_CASE_     )
													model.to(SCREAMING_SNAKE_CASE_     )
													model.eval()

													__UpperCamelCase       					=     model(SCREAMING_SNAKE_CASE_     )

													self.parent.assertEqual(result.last_hidden_state.shape       ,		(self.batch_size, self.seq_length, self.hidden_size)     )
													self.parent.assertEqual(len(result.hidden_states     )       ,		config.num_hidden_layers + 1     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		*SCREAMING_SNAKE_CASE_     )->				Optional[int]:





													'''simple docstring'''






													__UpperCamelCase       					=     RwkvForCausalLM(SCREAMING_SNAKE_CASE_     )
													model.to(SCREAMING_SNAKE_CASE_     )
													model.eval()

													__UpperCamelCase       					=     model(SCREAMING_SNAKE_CASE_       ,		labels=SCREAMING_SNAKE_CASE_     )
													self.parent.assertEqual(result.loss.shape       ,		()     )
													self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.seq_length, self.vocab_size)     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		*SCREAMING_SNAKE_CASE_     )->				Optional[int]:





													'''simple docstring'''






													__UpperCamelCase       					=     RwkvModel(config=SCREAMING_SNAKE_CASE_     )
													model.to(SCREAMING_SNAKE_CASE_     )
													model.eval()

													__UpperCamelCase       					=     model(SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     outputs.last_hidden_state

													__UpperCamelCase       					=     model(input_ids[:, :2]     )
													__UpperCamelCase       					=     outputs.last_hidden_state

													# Using the state computed on the first inputs, we will get the same output
													__UpperCamelCase       					=     model(input_ids[:, 2:]       ,		state=outputs.state     )
													__UpperCamelCase       					=     outputs.last_hidden_state

													self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]       ,		dim=1     )       ,		SCREAMING_SNAKE_CASE_       ,		atol=1E-5     )     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		*SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_=False     )->				List[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     RwkvForCausalLM(SCREAMING_SNAKE_CASE_     )
													model.to(SCREAMING_SNAKE_CASE_     )
													if gradient_checkpointing:
																			model.gradient_checkpointing_enable()

													__UpperCamelCase       					=     model(SCREAMING_SNAKE_CASE_       ,		labels=SCREAMING_SNAKE_CASE_     )
													self.parent.assertEqual(result.loss.shape       ,		()     )
													self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.seq_length, self.vocab_size)     )
													result.loss.backward()



							def       A__						(					self     )->				int:





													'''simple docstring'''






													__UpperCamelCase       					=     self.prepare_config_and_inputs()

													(
													    (
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,														
													)       					=     config_and_inputs

													__UpperCamelCase       					=     {'''input_ids''': input_ids}

													return config, inputs_dict






@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0				,						reason='See https://github.com/huggingface/transformers/pull/24204'      )
@require_torch
class     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						unittest.TestCase      ):

							"""simple docstring"""






							_snake_case       						=					(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
							_snake_case       						=					(
							    {'feature-extraction': RwkvModel, 'text-generation': RwkvForCausalLM} if is_torch_available() else {}
							)
							# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
							_snake_case       						=					False
							_snake_case       						=					False
							_snake_case       						=					False
							_snake_case       						=					False
							_snake_case       						=					False  # Rwkv does not support head masking



							def       A__						(					self     )->				List[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     RwkvModelTester(self     )
													__UpperCamelCase       					=     ConfigTester(
													    self       ,		config_class=SCREAMING_SNAKE_CASE_       ,		n_embd=37       ,		common_properties=['''hidden_size''', '''num_hidden_layers''']     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_=None     )->				List[str]:





													'''simple docstring'''






													if isinstance(SCREAMING_SNAKE_CASE_       ,		torch.Tensor     ):
																			__UpperCamelCase					,							__UpperCamelCase       					=     member.max().item(), member.min().item()
													elif isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     ) or isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     ):
																			__UpperCamelCase					,							__UpperCamelCase       					=     max(SCREAMING_SNAKE_CASE_     ), min(SCREAMING_SNAKE_CASE_     )

													if not isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     ):
																			raise TypeError('''container should be a list or tuple'''     )
													elif len(SCREAMING_SNAKE_CASE_     ) != 2:
																			raise ValueError('''container should have 2 elements'''     )

													__UpperCamelCase					,							__UpperCamelCase       					=     container

													__UpperCamelCase       					=     (min_value >= expected_min) and (max_value <= expected_max)

													if not is_inside_interval:
																			__UpperCamelCase       					=     '''%s not found in %s''' % (safe_repr(SCREAMING_SNAKE_CASE_     ), safe_repr(SCREAMING_SNAKE_CASE_     ))
																			self.fail(self._formatMessage(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )     )



							def       A__						(					self     )->				List[str]:





													'''simple docstring'''






													self.config_tester.run_common_tests()



							def       A__						(					self     )->				str:





													'''simple docstring'''






													__UpperCamelCase       					=     self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_rwkv_model(*SCREAMING_SNAKE_CASE_     )



							def       A__						(					self     )->				List[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_causl_lm(*SCREAMING_SNAKE_CASE_     )



							def       A__						(					self     )->				Tuple:





													'''simple docstring'''






													__UpperCamelCase       					=     self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_state_equivalency(*SCREAMING_SNAKE_CASE_     )



							def       A__						(					self     )->				Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase					,							__UpperCamelCase       					=     self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			__UpperCamelCase       					=     model_class(config=SCREAMING_SNAKE_CASE_     )
																			for name, param in model.named_parameters():
																									if "time_decay" in name:
																															if param.requires_grad:
																																					self.assertTrue(param.data.max().item() == 3.0     )
																																					self.assertTrue(param.data.min().item() == -5.0     )
																									elif "time_first" in name:
																															if param.requires_grad:
																																					# check if it's a ones like
																																					self.assertTrue(torch.allclose(param.data       ,		torch.ones_like(param.data     )       ,		atol=1E-5       ,		rtol=1E-5     )     )
																									elif any(x in name for x in ['''time_mix_key''', '''time_mix_receptance''']     ):
																															if param.requires_grad:
																																					self.assertInterval(
																																					    param.data       ,		[0.0, 1.0]       ,		msg=F"Parameter {name} of model {model_class} seems not properly initialized"       ,		)
																									elif "time_mix_value" in name:
																															if param.requires_grad:
																																					self.assertInterval(
																																					    param.data       ,		[0.0, 1.3]       ,		msg=F"Parameter {name} of model {model_class} seems not properly initialized"       ,		)



							def       A__						(					self     )->				str:





													'''simple docstring'''






													__UpperCamelCase					,							__UpperCamelCase       					=     self.model_tester.prepare_config_and_inputs_for_common()
													__UpperCamelCase       					=     True

													__UpperCamelCase       					=     getattr(self.model_tester       ,		'''seq_length'''       ,		SCREAMING_SNAKE_CASE_     )

													for model_class in self.all_model_classes:
																			__UpperCamelCase       					=     True
																			__UpperCamelCase       					=     False
																			__UpperCamelCase       					=     True
																			__UpperCamelCase       					=     model_class(SCREAMING_SNAKE_CASE_     )
																			model.to(SCREAMING_SNAKE_CASE_     )
																			model.eval()

																			__UpperCamelCase       					=     self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )
																			__UpperCamelCase       					=     inputs['''input_ids'''].shape[0]
																			with torch.no_grad():
																									__UpperCamelCase       					=     model(**SCREAMING_SNAKE_CASE_     )
																			__UpperCamelCase       					=     outputs.attentions
																			self.assertEqual(len(SCREAMING_SNAKE_CASE_     )       ,		self.model_tester.num_hidden_layers     )

																			# check that output_attentions also work using config
																			del inputs_dict["output_attentions"]
																			__UpperCamelCase       					=     True
																			__UpperCamelCase       					=     model_class(SCREAMING_SNAKE_CASE_     )
																			model.to(SCREAMING_SNAKE_CASE_     )
																			model.eval()

																			__UpperCamelCase       					=     self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )
																			__UpperCamelCase       					=     inputs['''input_ids'''].shape[0]
																			with torch.no_grad():
																									__UpperCamelCase       					=     model(**SCREAMING_SNAKE_CASE_     )
																			__UpperCamelCase       					=     outputs.attentions
																			self.assertEqual(len(SCREAMING_SNAKE_CASE_     )       ,		self.model_tester.num_hidden_layers     )

																			self.assertListEqual(
																			    list(attentions[0].shape[-3:]     )       ,		[batch_size, seq_len, config.hidden_size]       ,		)
																			__UpperCamelCase       					=     len(SCREAMING_SNAKE_CASE_     )

																			# Check attention is always last and order is fine
																			__UpperCamelCase       					=     True
																			__UpperCamelCase       					=     True
																			__UpperCamelCase       					=     model_class(SCREAMING_SNAKE_CASE_     )
																			model.to(SCREAMING_SNAKE_CASE_     )
																			model.eval()

																			__UpperCamelCase       					=     self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )
																			__UpperCamelCase       					=     inputs['''input_ids'''].shape[0]
																			with torch.no_grad():
																									__UpperCamelCase       					=     model(**SCREAMING_SNAKE_CASE_     )

																			__UpperCamelCase       					=     1
																			self.assertEqual(out_len + added_hidden_states       ,		len(SCREAMING_SNAKE_CASE_     )     )

																			__UpperCamelCase       					=     outputs.attentions

																			self.assertEqual(len(SCREAMING_SNAKE_CASE_     )       ,		self.model_tester.num_hidden_layers     )
																			self.assertListEqual(
																			    list(self_attentions[0].shape[-3:]     )       ,		[batch_size, seq_len, config.hidden_size]       ,		)



							@slow
							def       A__						(					self     )->				int:





													'''simple docstring'''






													for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			__UpperCamelCase       					=     RwkvModel.from_pretrained(SCREAMING_SNAKE_CASE_     )
																			self.assertIsNotNone(SCREAMING_SNAKE_CASE_     )






@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0				,						reason='See https://github.com/huggingface/transformers/pull/24204'      )
@slow
class     SCREAMING_SNAKE_CASE__      (							unittest.TestCase      ):

							"""simple docstring"""






							def       A__						(					self     )->				Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     '''RWKV/rwkv-4-169m-pile'''
													__UpperCamelCase       					=     AutoTokenizer.from_pretrained(self.model_id     )



							def       A__						(					self     )->				Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     '''Hello my name is Jasmine and I am a newbie to the'''
													__UpperCamelCase       					=     RwkvForCausalLM.from_pretrained(self.model_id     ).to(SCREAMING_SNAKE_CASE_     )

													__UpperCamelCase       					=     self.tokenizer('''Hello my name is'''       ,		return_tensors='''pt'''     ).input_ids.to(SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     model.generate(SCREAMING_SNAKE_CASE_       ,		max_new_tokens=10     )
													__UpperCamelCase       					=     self.tokenizer.decode(output[0].tolist()     )

													self.assertEqual(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )



							def       A__						(					self     )->				List[str]:





													'''simple docstring'''






													__UpperCamelCase       					=     '''Hello my name is Jasmine and I am a newbie to the'''

													__UpperCamelCase       					=     self.tokenizer('''Hello my name is'''       ,		return_tensors='''pt'''     ).input_ids.to(SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     RwkvForCausalLM.from_pretrained(self.model_id       ,		torch_dtype=torch.bfloataa     ).to(SCREAMING_SNAKE_CASE_     )

													__UpperCamelCase       					=     model.generate(SCREAMING_SNAKE_CASE_       ,		max_new_tokens=10     )
													__UpperCamelCase       					=     self.tokenizer.decode(output[0].tolist()     )

													self.assertEqual(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )






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 		A_	(							snake_case							:				Union[str, Any]	, snake_case							:				str	, snake_case							:				Dict	, snake_case							:				Optional[Any]=None	) ->      Optional[Any]:



						'''simple docstring'''

						__UpperCamelCase       					=     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.
												__UpperCamelCase       					=     weight.view(torch.intaa	)
												__UpperCamelCase       					=     '''bfloat16'''
						__UpperCamelCase       					=     weight.cpu().numpy()
						__UpperCamelCase       					=     os.path.join(snake_case	, f"{weight_name}.dat"	)
						if index is not None:
												if dtype is None:
																		__UpperCamelCase       					=     str(array.dtype	)
												__UpperCamelCase       					=     {'''dtype''': dtype, '''shape''': list(array.shape	)}
						if array.ndim == 0:
												__UpperCamelCase       					=     array[None]
						__UpperCamelCase       					=     np.memmap(snake_case	, dtype=array.dtype	, mode='''w+'''	, shape=array.shape	)
						__UpperCamelCase       					=     array[:]
						file_array.flush()
						return index




def 		A_	(							snake_case							:				Tuple	, snake_case							:				Union[str, Any]	) ->      Dict:



						'''simple docstring'''

						__UpperCamelCase       					=     tuple(weight_info['''shape''']	)
						if shape == ():
												# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
												__UpperCamelCase       					=     (1,)

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

						__UpperCamelCase       					=     np.memmap(snake_case	, dtype=snake_case	, shape=snake_case	, mode='''r'''	)

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

						return weight




def 		A_	(							snake_case							:				Dict	, snake_case							:				int	) ->      str:



						'''simple docstring'''

						if index is None or len(snake_case	) == 0:
												# Nothing to save
												return

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

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




def 		A_	(							snake_case							:				Union[str, os.PathLike]	, snake_case							:				Dict[str, torch.Tensor]	) ->      Dict:



						'''simple docstring'''

						os.makedirs(snake_case	, exist_ok=snake_case	)
						__UpperCamelCase       					=     {}
						for name, parameter in state_dict.items():
												__UpperCamelCase       					=     offload_weight(snake_case	, snake_case	, snake_case	, index=snake_case	)

						# Update index
						save_offload_index(snake_case	, snake_case	)






class     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_      ):

							"""simple docstring"""






							def __init__(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )->				str:





													'''simple docstring'''






													__UpperCamelCase       					=     dataset
													__UpperCamelCase       					=     prefix



							def __getitem__(					self       ,		SCREAMING_SNAKE_CASE_     )->				Union[str, Any]:





													'''simple docstring'''






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



							def __iter__(					self     )->				Union[str, Any]:





													'''simple docstring'''






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



							def __len__(					self     )->				str:





													'''simple docstring'''






													return len(self.dataset     )






class     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_      ):

							"""simple docstring"""






							def __init__(					self       ,		SCREAMING_SNAKE_CASE_ = None       ,		SCREAMING_SNAKE_CASE_ = None       ,		SCREAMING_SNAKE_CASE_ = None       ,		SCREAMING_SNAKE_CASE_=None       ,		)->				Tuple:





													'''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.'''     )

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



							def __getitem__(					self       ,		SCREAMING_SNAKE_CASE_     )->				Tuple:





													'''simple docstring'''






													if key in self.state_dict:
																			return self.state_dict[key]
													__UpperCamelCase       					=     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

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

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

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



							def __iter__(					self     )->				Dict:





													'''simple docstring'''






													return iter(self.all_keys     )



							def __len__(					self     )->				Dict:





													'''simple docstring'''






													return len(self.all_keys     )





def 		A_	(							snake_case							:				Dict[str, torch.Tensor]	, snake_case							:				List[str]	) ->      Tuple:



						'''simple docstring'''

						__UpperCamelCase       					=     {}
						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




