



import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_lowerCAmelCase      :    Dict						       =    namedtuple("ExampleNamedTuple", "a b c")






class 			__magic_name__					(    unittest.TestCase   ):





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




						'''simple docstring'''
						__a						   =torch.randn(5      ,						2      )
						__a						   =torch.device('cuda'      ) if torch.cuda.is_available() else torch.device('cpu'      )

						__a						   =send_to_device(_lowerCAmelCase      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta.cpu()      ,						_lowerCAmelCase      )      )

						__a						   =send_to_device((tensor, [tensor, tensor], 1)      ,						_lowerCAmelCase      )
						self.assertIsInstance(_lowerCAmelCase      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta[0].cpu()      ,						_lowerCAmelCase      )      )
						self.assertIsInstance(resulta[1]      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta[1][0].cpu()      ,						_lowerCAmelCase      )      )
						self.assertTrue(torch.equal(resulta[1][1].cpu()      ,						_lowerCAmelCase      )      )
						self.assertEqual(resulta[2]      ,						1      )

						__a						   =send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}      ,						_lowerCAmelCase      )
						self.assertIsInstance(_lowerCAmelCase      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta['a'].cpu()      ,						_lowerCAmelCase      )      )
						self.assertIsInstance(resulta['b']      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta['b'][0].cpu()      ,						_lowerCAmelCase      )      )
						self.assertTrue(torch.equal(resulta['b'][1].cpu()      ,						_lowerCAmelCase      )      )
						self.assertEqual(resulta['c']      ,						1      )

						__a						   =send_to_device(ExampleNamedTuple(a=_lowerCAmelCase      ,						b=[tensor, tensor]      ,						c=1      )      ,						_lowerCAmelCase      )
						self.assertIsInstance(_lowerCAmelCase      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta.a.cpu()      ,						_lowerCAmelCase      )      )
						self.assertIsInstance(resulta.b      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta.b[0].cpu()      ,						_lowerCAmelCase      )      )
						self.assertTrue(torch.equal(resulta.b[1].cpu()      ,						_lowerCAmelCase      )      )
						self.assertEqual(resulta.c      ,						1      )

						__a						   =send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1}      )      ,						_lowerCAmelCase      )
						self.assertIsInstance(_lowerCAmelCase      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta['a'].cpu()      ,						_lowerCAmelCase      )      )
						self.assertIsInstance(resulta['b']      ,						_lowerCAmelCase      )
						self.assertTrue(torch.equal(resulta['b'][0].cpu()      ,						_lowerCAmelCase      )      )
						self.assertTrue(torch.equal(resulta['b'][1].cpu()      ,						_lowerCAmelCase      )      )
						self.assertEqual(resulta['c']      ,						1      )





			def 	__magic_name__					(	self      )     ->				int:




						'''simple docstring'''
						with self.assertRaises(_lowerCAmelCase      ) as cm:
									__a						   =recursively_apply(torch.tensor      ,						(torch.tensor(1      ), 1)      ,						error_on_other_type=_lowerCAmelCase      )
						self.assertEqual(
						    str(cm.exception      )      ,						'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'      ,						)





			def 	__magic_name__					(	self      )     ->				Any:




						'''simple docstring'''
						__a						   =torch.tensor([1, 2, 3, 4, 5]      )
						self.assertEqual(listify(_lowerCAmelCase      )      ,						[1, 2, 3, 4, 5]      )

						__a						   =torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]      )
						self.assertEqual(listify(_lowerCAmelCase      )      ,						[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]      )

						__a						   =torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]      )
						self.assertEqual(
						    listify(_lowerCAmelCase      )      ,						[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]      )





			def 	__magic_name__					(	self      )     ->				str:




						'''simple docstring'''
						with patch_environment(aa=1      ,						BB=2      ):
									self.assertEqual(os.environ.get('AA'      )      ,						'1'      )
									self.assertEqual(os.environ.get('BB'      )      ,						'2'      )

						self.assertNotIn('AA'      ,						os.environ      )
						self.assertNotIn('BB'      ,						os.environ      )





			def 	__magic_name__					(	self      )     ->				Tuple:




						'''simple docstring'''
						__a						   =RegressionModel()
						__a						   =model.forward
						__a						   =convert_outputs_to_fpaa(model.forward      )
						__a						   =extract_model_from_parallel(_lowerCAmelCase      ,						keep_fpaa_wrapper=_lowerCAmelCase      )
						__a						   =pickle.dumps(_lowerCAmelCase      )





			@require_cuda
			def 	__magic_name__					(	self      )     ->				Union[str, Any]:




						'''simple docstring'''
						__a						   =RegressionModel()
						__a						   =model.forward
						__a						   =torch.cuda.amp.autocast(dtype=torch.floataa      )(model.forward      )
						__a						   =convert_outputs_to_fpaa(model.forward      )
						__a						   =extract_model_from_parallel(_lowerCAmelCase      ,						keep_fpaa_wrapper=_lowerCAmelCase      )
						__a						   =pickle.dumps(_lowerCAmelCase      )





			@require_cuda
			@require_torch_min_version(version='2.0'      )
			def 	__magic_name__					(	self      )     ->				str:




						'''simple docstring'''
						__a						   =RegressionModel()
						__a						   =model.forward
						__a						   =torch.cuda.amp.autocast(dtype=torch.floataa      )(model.forward      )
						__a						   =convert_outputs_to_fpaa(model.forward      )
						__a						   =torch.compile(model.forward      ,						backend='inductor'      )
						__a						   =torch.randn(4      ,						10      ).cuda()
						__a						   =model(_lowerCAmelCase      )





			def 	__magic_name__					(	self      )     ->				Any:




						'''simple docstring'''
						__a						   =RegressionModel()
						# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
						__a						   =torch.nn.parallel.DataParallel(_lowerCAmelCase      )
						__a						   =extract_model_from_parallel(_lowerCAmelCase      )

						self.assertEqual(_lowerCAmelCase      ,						_lowerCAmelCase      )





			@require_torch_min_version(version='2.0'      )
			def 	__magic_name__					(	self      )     ->				Optional[int]:




						'''simple docstring'''
						__a						   =RegressionModel()
						__a						   =torch.compile(_lowerCAmelCase      )

						# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
						__a						   =torch.nn.parallel.DataParallel(_lowerCAmelCase      )
						__a						   =torch.compile(_lowerCAmelCase      )
						__a						   =extract_model_from_parallel(_lowerCAmelCase      )

						self.assertEqual(compiled_model._orig_mod      ,						compiled_model_unwrapped._orig_mod      )





			def 	__magic_name__					(	self      )     ->				Dict:




						'''simple docstring'''
						self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3]      )]      )      ,						torch.device('cpu'      )      )
						self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3]      )}      )      ,						torch.device('cpu'      )      )
						self.assertIsNone(find_device([1, 'a']      )      )









from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_lowerCAmelCase      :    str						       =    logging.get_logger(__name__)

_lowerCAmelCase      :    List[str]						       =    {
    "facebook/mbart-large-cc25": "https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}






class 			__magic_name__					(    lowerCAmelCase_   ):
			SCREAMING_SNAKE_CASE							     =		'mbart'
			SCREAMING_SNAKE_CASE							     =		['past_key_values']
			SCREAMING_SNAKE_CASE							     =		{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}





			def __init__(	self      ,						__snake_case=5_0265      ,						__snake_case=1024      ,						__snake_case=12      ,						__snake_case=4096      ,						__snake_case=16      ,						__snake_case=12      ,						__snake_case=4096      ,						__snake_case=16      ,						__snake_case=0.0      ,						__snake_case=0.0      ,						__snake_case=True      ,						__snake_case=True      ,						__snake_case="gelu"      ,						__snake_case=1024      ,						__snake_case=0.1      ,						__snake_case=0.0      ,						__snake_case=0.0      ,						__snake_case=0.02      ,						__snake_case=0.0      ,						__snake_case=False      ,						__snake_case=1      ,						__snake_case=0      ,						__snake_case=2      ,						__snake_case=2      ,						**__snake_case      ,						)     ->				Optional[Any]:




						'''simple docstring'''
						__a						   =vocab_size
						__a						   =max_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						   =scale_embedding  # scale factor will be sqrt(d_model) if True
						super().__init__(
						    pad_token_id=__snake_case      ,						bos_token_id=__snake_case      ,						eos_token_id=__snake_case      ,						is_encoder_decoder=__snake_case      ,						forced_eos_token_id=__snake_case      ,						**__snake_case      ,						)






class 			__magic_name__					(    lowerCAmelCase_   ):





			@property
			def 	__magic_name__					(	self      )     ->				Mapping[str, Mapping[int, str]]:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =OrderedDict(
									    [
									        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
									        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
									    ]      )

									if self.use_past:
												__a						   ={0: 'batch'}
												__a						   ={0: 'batch', 1: 'past_decoder_sequence + sequence'}
									else:
												__a						   ={0: 'batch', 1: 'decoder_sequence'}
												__a						   ={0: 'batch', 1: 'decoder_sequence'}

									if self.use_past:
												self.fill_with_past_key_values_(__snake_case      ,						direction='inputs'      )
						elif self.task == "causal-lm":
									# TODO: figure this case out.
									__a						   =OrderedDict(
									    [
									        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
									        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
									    ]      )
									if self.use_past:
												__a   ,					__a						   =self.num_layers
												for i in range(__snake_case      ):
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
						else:
									__a						   =OrderedDict(
									    [
									        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
									        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
									        ('decoder_input_ids', {0: 'batch', 1: 'decoder_sequence'}),
									        ('decoder_attention_mask', {0: 'batch', 1: 'decoder_sequence'}),
									    ]      )

						return common_inputs





			@property
			def 	__magic_name__					(	self      )     ->				Mapping[str, Mapping[int, str]]:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =super().outputs
						else:
									__a						   =super(__snake_case      ,						self      ).outputs
									if self.use_past:
												__a   ,					__a						   =self.num_layers
												for i in range(__snake_case      ):
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
						return common_outputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )

						# Generate decoder inputs
						__a						   =seq_length if not self.use_past else 1
						__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )
						__a						   ={f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
						__a						   =dict(**__snake_case      ,						**__snake_case      )

						if self.use_past:
									if not is_torch_available():
												raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'      )
									else:
												import torch
									__a   ,					__a						   =common_inputs['input_ids'].shape
									__a						   =common_inputs['decoder_input_ids'].shape[1]
									__a   ,					__a						   =self.num_attention_heads
									__a						   =(
									    batch,
									    num_encoder_attention_heads,
									    encoder_seq_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)
									__a						   =decoder_seq_length + 3
									__a						   =(
									    batch,
									    num_decoder_attention_heads,
									    decoder_past_length,
									    self._config.hidden_size // num_decoder_attention_heads,
									)

									__a						   =torch.cat(
									    [common_inputs['decoder_attention_mask'], torch.ones(__snake_case      ,						__snake_case      )]      ,						dim=1      )

									__a						   =[]
									# If the number of encoder and decoder layers are present in the model configuration, both are considered
									__a   ,					__a						   =self.num_layers
									__a						   =min(__snake_case      ,						__snake_case      )
									__a						   =max(__snake_case      ,						__snake_case      ) - min_num_layers
									__a						   ='encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

									for _ in range(__snake_case      ):
												common_inputs["past_key_values"].append(
												    (
												        torch.zeros(__snake_case      ),
												        torch.zeros(__snake_case      ),
												        torch.zeros(__snake_case      ),
												        torch.zeros(__snake_case      ),
												    )      )
									# TODO: test this.
									__a						   =encoder_shape if remaining_side_name == 'encoder' else decoder_shape
									for _ in range(__snake_case      ,						__snake_case      ):
												common_inputs["past_key_values"].append((torch.zeros(__snake_case      ), torch.zeros(__snake_case      ))      )
						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )

						if self.use_past:
									if not is_torch_available():
												raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'      )
									else:
												import torch
									__a   ,					__a						   =common_inputs['input_ids'].shape
									# Not using the same length for past_key_values
									__a						   =seqlen + 2
									__a   ,					__a						   =self.num_layers
									__a   ,					__a						   =self.num_attention_heads
									__a						   =(
									    batch,
									    num_encoder_attention_heads,
									    past_key_values_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)

									__a						   =common_inputs['attention_mask'].dtype
									__a						   =torch.cat(
									    [common_inputs['attention_mask'], torch.ones(__snake_case      ,						__snake_case      ,						dtype=__snake_case      )]      ,						dim=1      )
									__a						   =[
									    (torch.zeros(__snake_case      ), torch.zeros(__snake_case      )) for _ in range(__snake_case      )
									]
						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						# Copied from OnnxConfig.generate_dummy_inputs
						# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
						# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
						__a						   =compute_effective_axis_dimension(
						    __snake_case      ,						fixed_dimension=OnnxConfig.default_fixed_batch      ,						num_token_to_add=0      )

						# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
						__a						   =tokenizer.num_special_tokens_to_add(__snake_case      )
						__a						   =compute_effective_axis_dimension(
						    __snake_case      ,						fixed_dimension=OnnxConfig.default_fixed_sequence      ,						num_token_to_add=__snake_case      )

						# Generate dummy inputs according to compute batch and sequence
						__a						   =[' '.join([tokenizer.unk_token]      ) * seq_length] * batch_size
						__a						   =dict(tokenizer(__snake_case      ,						return_tensors=__snake_case      )      )
						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =self._generate_dummy_inputs_for_default_and_seqaseq_lm(
									    __snake_case      ,						batch_size=__snake_case      ,						seq_length=__snake_case      ,						is_pair=__snake_case      ,						framework=__snake_case      )

						elif self.task == "causal-lm":
									__a						   =self._generate_dummy_inputs_for_causal_lm(
									    __snake_case      ,						batch_size=__snake_case      ,						seq_length=__snake_case      ,						is_pair=__snake_case      ,						framework=__snake_case      )
						else:
									__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    __snake_case      ,						batch_size=__snake_case      ,						seq_length=__snake_case      ,						is_pair=__snake_case      ,						framework=__snake_case      )

						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )     ->				str:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =super()._flatten_past_key_values_(__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )
						else:
									__a						   =super(__snake_case      ,						self      )._flatten_past_key_values_(
									    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )






