code: 
  
  
  from __future__ import annotations
  
  import unittest
  import warnings
  
  from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
  from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
  from transformers.utils import cached_property
  
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
  from ...test_pipeline_mixin import PipelineTesterMixin
  
  
  if is_tf_available():
  			import tensorflow as tf
  
  			from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel
  
  @require_tf
  class 						snake_case			:
  			'''simple docstring'''
  
  			UpperCAmelCase       :       List[Any]       = MarianConfig
  			UpperCAmelCase       :       Dict       = {}
  			UpperCAmelCase       :       Optional[int]       = "gelu"
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     List[str]							,							lowerCAmelCase_     :     Optional[int]=13							,							lowerCAmelCase_     :     List[str]=7							,							lowerCAmelCase_     :     int=True							,							lowerCAmelCase_     :     Optional[Any]=False							,							lowerCAmelCase_     :     List[Any]=99							,							lowerCAmelCase_     :     Optional[Any]=32							,							lowerCAmelCase_     :     List[Any]=2							,							lowerCAmelCase_     :     str=4							,							lowerCAmelCase_     :     Dict=37							,							lowerCAmelCase_     :     int=0.1							,							lowerCAmelCase_     :     str=0.1							,							lowerCAmelCase_     :     List[str]=20							,							lowerCAmelCase_     :     str=2							,							lowerCAmelCase_     :     Any=1							,							lowerCAmelCase_     :     List[str]=0							,							)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       parent
  								SCREAMING_SNAKE_CASE_	      =       batch_size
  								SCREAMING_SNAKE_CASE_	      =       seq_length
  								SCREAMING_SNAKE_CASE_	      =       is_training
  								SCREAMING_SNAKE_CASE_	      =       use_labels
  								SCREAMING_SNAKE_CASE_	      =       vocab_size
  								SCREAMING_SNAKE_CASE_	      =       hidden_size
  								SCREAMING_SNAKE_CASE_	      =       num_hidden_layers
  								SCREAMING_SNAKE_CASE_	      =       num_attention_heads
  								SCREAMING_SNAKE_CASE_	      =       intermediate_size
  
  								SCREAMING_SNAKE_CASE_	      =       hidden_dropout_prob
  								SCREAMING_SNAKE_CASE_	      =       attention_probs_dropout_prob
  								SCREAMING_SNAKE_CASE_	      =       max_position_embeddings
  								SCREAMING_SNAKE_CASE_	      =       eos_token_id
  								SCREAMING_SNAKE_CASE_	      =       pad_token_id
  								SCREAMING_SNAKE_CASE_	      =       bos_token_id
  
  			def _lowercase (	self     :     Any					)					->					List[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       ids_tensor([self.batch_size, self.seq_length - 1]							,							self.vocab_size					)
  								SCREAMING_SNAKE_CASE_	      =       tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size					)							,							1					)
  								SCREAMING_SNAKE_CASE_	      =       tf.concat([input_ids, eos_tensor]							,							axis=1					)
  
  								SCREAMING_SNAKE_CASE_	      =       ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size					)
  
  								SCREAMING_SNAKE_CASE_	      =       self.config_cls(
  								    vocab_size=self.vocab_size							,							d_model=self.hidden_size							,							encoder_layers=self.num_hidden_layers							,							decoder_layers=self.num_hidden_layers							,							encoder_attention_heads=self.num_attention_heads							,							decoder_attention_heads=self.num_attention_heads							,							encoder_ffn_dim=self.intermediate_size							,							decoder_ffn_dim=self.intermediate_size							,							dropout=self.hidden_dropout_prob							,							attention_dropout=self.attention_probs_dropout_prob							,							max_position_embeddings=self.max_position_embeddings							,							eos_token_ids=[2]							,							bos_token_id=self.bos_token_id							,							pad_token_id=self.pad_token_id							,							decoder_start_token_id=self.pad_token_id							,							**self.config_updates							,							)
  								SCREAMING_SNAKE_CASE_	      =       prepare_marian_inputs_dict(UpperCAmelCase_							,							UpperCAmelCase_							,							UpperCAmelCase_					)
  								return config, inputs_dict
  
  
  
  
  			def _lowercase (	self     :     Optional[int]							,							lowerCAmelCase_     :     str							,							lowerCAmelCase_     :     Optional[Any]					)					->					Optional[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       TFMarianModel(config=UpperCAmelCase_					).get_decoder()
  								SCREAMING_SNAKE_CASE_	      =       inputs_dict['input_ids']
  
  								SCREAMING_SNAKE_CASE_	      =       input_ids[:1, :]
  								SCREAMING_SNAKE_CASE_	      =       inputs_dict['attention_mask'][:1, :]
  								SCREAMING_SNAKE_CASE_	      =       inputs_dict['head_mask']
  								SCREAMING_SNAKE_CASE_	      =       1
  
  								# first forward pass
  								SCREAMING_SNAKE_CASE_	      =       model(UpperCAmelCase_							,							attention_mask=UpperCAmelCase_							,							head_mask=UpperCAmelCase_							,							use_cache=UpperCAmelCase_					)
  
  								SCREAMING_SNAKE_CASE_	      =       outputs.to_tuple()
  
  								# create hypothetical next token and extent to next_input_ids
  								SCREAMING_SNAKE_CASE_	      =       ids_tensor((self.batch_size, 3)							,							config.vocab_size					)
  								SCREAMING_SNAKE_CASE_	      =       tf.cast(ids_tensor((self.batch_size, 3)							,							2					)							,							tf.inta					)
  
  								# append to next input_ids and
  								SCREAMING_SNAKE_CASE_	      =       tf.concat([input_ids, next_tokens]							,							axis=-1					)
  								SCREAMING_SNAKE_CASE_	      =       tf.concat([attention_mask, next_attn_mask]							,							axis=-1					)
  
  								SCREAMING_SNAKE_CASE_	      =       model(UpperCAmelCase_							,							attention_mask=UpperCAmelCase_					)[0]
  								SCREAMING_SNAKE_CASE_	      =       model(UpperCAmelCase_							,							attention_mask=UpperCAmelCase_							,							past_key_values=UpperCAmelCase_					)[0]
  
  								self.parent.assertEqual(next_tokens.shape[1]							,							output_from_past.shape[1]					)
  
  								# select random slice
  								SCREAMING_SNAKE_CASE_	      =       int(ids_tensor((1,)							,							output_from_past.shape[-1]					)					)
  								SCREAMING_SNAKE_CASE_	      =       output_from_no_past[:, -3:, random_slice_idx]
  								SCREAMING_SNAKE_CASE_	      =       output_from_past[:, :, random_slice_idx]
  
  								# test that outputs are equal for slice
  								tf.debugging.assert_near(UpperCAmelCase_							,							UpperCAmelCase_							,							rtol=1e-3					)
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase=None				,UpperCAmelCase=None				,UpperCAmelCase=None				,UpperCAmelCase=None				,UpperCAmelCase=None				,)->		Dict:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					if attention_mask is None:
  										SCREAMING_SNAKE_CASE_	      =       tf.cast(tf.math.not_equal(_UpperCAmelCase				,config.pad_token_id      )				,tf.inta      )
  					if decoder_attention_mask is None:
  										SCREAMING_SNAKE_CASE_	      =       tf.concat(
  										    [
  										        tf.ones(decoder_input_ids[:, :1].shape				,dtype=tf.inta      ),
  										        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]				,config.pad_token_id      )				,tf.inta      ),
  										    ]				,axis=-1				,)
  					if head_mask is None:
  										SCREAMING_SNAKE_CASE_	      =       tf.ones((config.encoder_layers, config.encoder_attention_heads)      )
  					if decoder_head_mask is None:
  										SCREAMING_SNAKE_CASE_	      =       tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
  					if cross_attn_head_mask is None:
  										SCREAMING_SNAKE_CASE_	      =       tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
  					return {
  					    "input_ids": input_ids,
  					    "decoder_input_ids": decoder_input_ids,
  					    "attention_mask": attention_mask,
  					    "decoder_attention_mask": decoder_attention_mask,
  					    "head_mask": head_mask,
  					    "decoder_head_mask": decoder_head_mask,
  					    "cross_attn_head_mask": cross_attn_head_mask,
  					}
  
  @require_tf
  class 						snake_case			(       lowerCAmelCase__	, lowerCAmelCase__	, unittest.TestCase	):
  			'''simple docstring'''
  
  			UpperCAmelCase       :       Any       = (TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
  			UpperCAmelCase       :       Optional[Any]       = (TFMarianMTModel,) if is_tf_available() else ()
  			UpperCAmelCase       :       Union[str, Any]       = (
  			    {
  			        "conversational": TFMarianMTModel,
  			        "feature-extraction": TFMarianModel,
  			        "summarization": TFMarianMTModel,
  			        "text2text-generation": TFMarianMTModel,
  			        "translation": TFMarianMTModel,
  			    }
  			    if is_tf_available()
  			    else {}
  			)
  			UpperCAmelCase       :       int       = True
  			UpperCAmelCase       :       Union[str, Any]       = False
  			UpperCAmelCase       :       List[Any]       = False
  
  			def _lowercase (	self     :     str					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       TFMarianModelTester(self					)
  								SCREAMING_SNAKE_CASE_	      =       ConfigTester(self							,							config_class=UpperCAmelCase_					)
  
  			def _lowercase (	self     :     int					)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								self.config_tester.run_common_tests()
  
  
  
  
  			def _lowercase (	self     :     List[str]					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       self.model_tester.prepare_config_and_inputs_for_common()
  								self.model_tester.check_decoder_model_past_large_inputs(*UpperCAmelCase_					)
  
  @require_tf
  class 						snake_case			(       unittest.TestCase	):
  			'''simple docstring'''
  
  			UpperCAmelCase       :       Optional[int]       = 1000  # show more chars for failing integration tests
  
  			@classmethod
  			def _lowercase (	cls     :     Dict					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       F'''Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}'''
  								return cls
  
  			@cached_property
  			def _lowercase (	self     :     List[str]					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return AutoTokenizer.from_pretrained(self.model_name					)
  
  			@property
  			def _lowercase (	self     :     int					)					->					List[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.tokenizer.eos_token_id
  
  			@cached_property
  			def _lowercase (	self     :     Optional[int]					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								warnings.simplefilter('''error'''					)
  								SCREAMING_SNAKE_CASE_	      =       TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name					)
  								assert isinstance(UpperCAmelCase_							,							UpperCAmelCase_					)
  								SCREAMING_SNAKE_CASE_	      =       model.config
  								self.assertListEqual(c.bad_words_ids							,							[[c.pad_token_id]]					)
  								self.assertEqual(c.max_length							,							512					)
  								self.assertEqual(c.decoder_start_token_id							,							c.pad_token_id					)
  								return model
  
  			def _lowercase (	self     :     int							,							**lowerCAmelCase_     :     Optional[Any]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       self.translate_src_text(**UpperCAmelCase_					)
  								self.assertListEqual(self.expected_text							,							UpperCAmelCase_					)
  
  
  
  
  			def _lowercase (	self     :     List[str]							,							**lowerCAmelCase_     :     Optional[Any]					)					->					Any:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       self.tokenizer(self.src_text							,							**UpperCAmelCase_							,							padding=UpperCAmelCase_							,							return_tensors='''tf'''					)
  								SCREAMING_SNAKE_CASE_	      =       self.model.generate(
  								    model_inputs.input_ids							,							attention_mask=model_inputs.attention_mask							,							num_beams=2							,							max_length=128					)
  								SCREAMING_SNAKE_CASE_	      =       self.tokenizer.batch_decode(generated_ids.numpy()							,							skip_special_tokens=UpperCAmelCase_					)
  								return generated_words
  
  @require_sentencepiece
  @require_tokenizers
  @require_tf
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			UpperCAmelCase       :       Optional[Any]       = "mt"
  			UpperCAmelCase       :       Dict       = "en"
  			UpperCAmelCase       :       Union[str, Any]       = ["Billi messu b'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem."]
  			UpperCAmelCase       :       Dict       = ["Touching gently, Jesus healed a man who was affected by the sad disease of leprosy."]
  
  			@unittest.skip('''Skipping until #12647 is resolved.'''					)
  			@slow
  			def _lowercase (	self     :     Tuple					)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								self._assert_generated_batch_equal_expected()
  
  @require_sentencepiece
  @require_tokenizers
  @require_tf
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			UpperCAmelCase       :       Optional[Any]       = "en"
  			UpperCAmelCase       :       int       = "zh"
  			UpperCAmelCase       :       int       = ["My name is Wolfgang and I live in Berlin"]
  			UpperCAmelCase       :       str       = ["我叫沃尔夫冈 我住在柏林"]
  
  			@unittest.skip('''Skipping until #12647 is resolved.'''					)
  			@slow
  			def _lowercase (	self     :     List[Any]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								self._assert_generated_batch_equal_expected()
  
  @require_sentencepiece
  @require_tokenizers
  @require_tf
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			UpperCAmelCase       :       str       = "en"
  			UpperCAmelCase       :       Tuple       = "ROMANCE"
  			UpperCAmelCase       :       Tuple       = [
  			    ">>fr<< Don't spend so much time watching TV.",
  			    ">>pt<< Your message has been sent.",
  			    ">>es<< He's two years older than me.",
  			]
  			UpperCAmelCase       :       Union[str, Any]       = [
  			    "Ne passez pas autant de temps à regarder la télé.",
  			    "A sua mensagem foi enviada.",
  			    "Es dos años más viejo que yo.",
  			]
  
  			@unittest.skip('''Skipping until #12647 is resolved.'''					)
  			@slow
  			def _lowercase (	self     :     Union[str, Any]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								self._assert_generated_batch_equal_expected()
  
  
  
  
  			@unittest.skip('''Skipping until #12647 is resolved.'''					)
  			@slow
  			def _lowercase (	self     :     Dict					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       TranslationPipeline(self.model							,							self.tokenizer							,							framework='''tf'''					)
  								SCREAMING_SNAKE_CASE_	      =       pipeline(self.src_text					)
  								self.assertEqual(self.expected_text							,							[x['''translation_text'''] for x in output]					)
  
  
  
style_context: 
  
  
  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				(     UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase=None      )->		Dict:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       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.
  										SCREAMING_SNAKE_CASE_	      =       weight.view(torch.intaa      )
  										SCREAMING_SNAKE_CASE_	      =       '''bfloat16'''
  					SCREAMING_SNAKE_CASE_	      =       weight.cpu().numpy()
  					SCREAMING_SNAKE_CASE_	      =       os.path.join(UpperCAmelCase				,f'''{weight_name}.dat'''      )
  					if index is not None:
  										if dtype is None:
  															SCREAMING_SNAKE_CASE_	      =       str(array.dtype      )
  										SCREAMING_SNAKE_CASE_	      =       {'''dtype''': dtype, '''shape''': list(array.shape      )}
  					if array.ndim == 0:
  										SCREAMING_SNAKE_CASE_	      =       array[None]
  					SCREAMING_SNAKE_CASE_	      =       np.memmap(UpperCAmelCase				,dtype=array.dtype				,mode='''w+'''				,shape=array.shape      )
  					SCREAMING_SNAKE_CASE_	      =       array[:]
  					file_array.flush()
  					return index
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		List[str]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       tuple(weight_info['''shape''']      )
  					if shape == ():
  										# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
  										SCREAMING_SNAKE_CASE_	      =       (1,)
  
  					SCREAMING_SNAKE_CASE_	      =       weight_info['''dtype''']
  					if dtype == "bfloat16":
  										# NumPy does not support bfloat16 so this was saved as a int16
  										SCREAMING_SNAKE_CASE_	      =       '''int16'''
  
  					SCREAMING_SNAKE_CASE_	      =       np.memmap(UpperCAmelCase				,dtype=UpperCAmelCase				,shape=UpperCAmelCase				,mode='''r'''      )
  
  					if len(weight_info['''shape''']      ) == 0:
  										SCREAMING_SNAKE_CASE_	      =       weight[0]
  					SCREAMING_SNAKE_CASE_	      =       torch.tensor(UpperCAmelCase      )
  					if weight_info["dtype"] == "bfloat16":
  										SCREAMING_SNAKE_CASE_	      =       weight.view(torch.bfloataa      )
  
  					return weight
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		str:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					if index is None or len(UpperCAmelCase      ) == 0:
  										# Nothing to save
  										return
  
  					SCREAMING_SNAKE_CASE_	      =       os.path.join(UpperCAmelCase				,'''index.json'''      )
  					if os.path.isfile(UpperCAmelCase      ):
  										with open(UpperCAmelCase				,'''r'''				,encoding='''utf-8'''      ) as f:
  															SCREAMING_SNAKE_CASE_	      =       json.load(UpperCAmelCase      )
  					else:
  										SCREAMING_SNAKE_CASE_	      =       {}
  					current_index.update(UpperCAmelCase      )
  
  					with open(UpperCAmelCase				,'''w'''				,encoding='''utf-8'''      ) as f:
  										json.dump(UpperCAmelCase				,UpperCAmelCase				,indent=2      )
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		Optional[Any]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					os.makedirs(UpperCAmelCase				,exist_ok=UpperCAmelCase      )
  					SCREAMING_SNAKE_CASE_	      =       {}
  					for name, parameter in state_dict.items():
  										SCREAMING_SNAKE_CASE_	      =       offload_weight(UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,index=UpperCAmelCase      )
  
  					# Update index
  					save_offload_index(UpperCAmelCase				,UpperCAmelCase      )
  
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     Mapping							,							lowerCAmelCase_     :     str					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       dataset
  								SCREAMING_SNAKE_CASE_	      =       prefix
  
  			def __getitem__(	self     :     Optional[int]							,							lowerCAmelCase_     :     Optional[int]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.dataset[F'''{self.prefix}{key}''']
  
  			def __iter__(	self     :     Tuple					)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return iter([key for key in self.dataset if key.startswith(self.prefix					)]					)
  
  
  
  
  			def __len__(	self     :     Dict					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return len(self.dataset					)
  
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     Dict[str, torch.Tensor] = None							,							lowerCAmelCase_     :     Optional[Union[str, os.PathLike]] = None							,							lowerCAmelCase_     :     Mapping = None							,							lowerCAmelCase_     :     int=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.'''					)
  
  								SCREAMING_SNAKE_CASE_	      =       {} if state_dict is None else state_dict
  								SCREAMING_SNAKE_CASE_	      =       save_folder
  								if index is None and save_folder is not None:
  													with open(os.path.join(lowerCAmelCase_							,							'''index.json'''					)					) as f:
  																		SCREAMING_SNAKE_CASE_	      =       json.load(lowerCAmelCase_					)
  								SCREAMING_SNAKE_CASE_	      =       {} if index is None else index
  								SCREAMING_SNAKE_CASE_	      =       list(self.state_dict.keys()					)
  								self.all_keys.extend([key for key in self.index if key not in self.all_keys]					)
  								SCREAMING_SNAKE_CASE_	      =       device
  
  			def __getitem__(	self     :     str							,							lowerCAmelCase_     :     str					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if key in self.state_dict:
  													return self.state_dict[key]
  								SCREAMING_SNAKE_CASE_	      =       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
  
  													SCREAMING_SNAKE_CASE_	      =       '''cpu''' if self.device is None else self.device
  													with safe_open(weight_info['''safetensors_file''']							,							framework='''pt'''							,							device=lowerCAmelCase_					) as f:
  																		SCREAMING_SNAKE_CASE_	      =       f.get_tensor(weight_info.get('''weight_name'''							,							lowerCAmelCase_					)					)
  
  													if "dtype" in weight_info:
  																		return tensor.to(getattr(lowerCAmelCase_							,							weight_info['''dtype''']					)					)
  													else:
  																		return tensor
  
  								SCREAMING_SNAKE_CASE_	      =       os.path.join(self.save_folder							,							F'''{key}.dat'''					)
  								return load_offloaded_weight(lowerCAmelCase_							,							lowerCAmelCase_					)
  
  			def __iter__(	self     :     str					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return iter(self.all_keys					)
  
  
  
  
  			def __len__(	self     :     Optional[Any]					)					->					Optional[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return len(self.all_keys					)
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		List[str]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       {}
  					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
  
  
  