





import argparse

import torch
from torch import nn

from transformers import SpeechaTextConfig, SpeechaTextForConditionalGeneration







def        __lowerCamelCase (							__a       :Any   )		->    Dict:






		"""simple docstring"""

		A__  				=       [
		    """encoder.version""",
		    """decoder.version""",
		    """model.encoder.version""",
		    """model.decoder.version""",
		    """decoder.output_projection.weight""",
		    """_float_tensor""",
		    """encoder.embed_positions._float_tensor""",
		    """decoder.embed_positions._float_tensor""",
		]
		for k in ignore_keys:
				state_dict.pop(lowerCAmelCase__   ,					lowerCAmelCase__   )







def        __lowerCamelCase (							__a       :Optional[Any]   )		->    Optional[int]:






		"""simple docstring"""

		A__  				=       list(s_dict.keys()   )
		for key in keys:
				if "transformer_layers" in key:
						A__  				=       s_dict.pop(lowerCAmelCase__   )
				elif "subsample" in key:
						A__  				=       s_dict.pop(lowerCAmelCase__   )







def        __lowerCamelCase (							__a       :Optional[Any]   )		->    List[Any]:






		"""simple docstring"""

		A__				,      A__  				=       emb.weight.shape
		A__  				=       nn.Linear(lowerCAmelCase__   ,					lowerCAmelCase__   ,					bias=lowerCAmelCase__   )
		A__  				=       emb.weight.data
		return lin_layer







def        __lowerCamelCase (							__a       :Dict   ,					__a       :List[str]   )		->    Optional[Any]:






		"""simple docstring"""

		A__  				=       torch.load(lowerCAmelCase__   ,					map_location="""cpu"""   )
		A__  				=       mam_aaa["""args"""]
		A__  				=       mam_aaa["""model"""]
		A__  				=       state_dict["""decoder.output_projection.weight"""]

		remove_ignore_keys_(lowerCAmelCase__   )
		rename_keys(lowerCAmelCase__   )

		A__  				=       state_dict["""decoder.embed_tokens.weight"""].shape[0]

		A__  				=       args.share_decoder_input_output_embed

		A__  				=       [int(lowerCAmelCase__   ) for i in args.conv_kernel_sizes.split(""","""   )]
		A__  				=       SpeechaTextConfig(
		    vocab_size=lowerCAmelCase__   ,					max_source_positions=args.max_source_positions   ,					max_target_positions=args.max_target_positions   ,					encoder_layers=args.encoder_layers   ,					decoder_layers=args.decoder_layers   ,					encoder_attention_heads=args.encoder_attention_heads   ,					decoder_attention_heads=args.decoder_attention_heads   ,					encoder_ffn_dim=args.encoder_ffn_embed_dim   ,					decoder_ffn_dim=args.decoder_ffn_embed_dim   ,					d_model=args.encoder_embed_dim   ,					dropout=args.dropout   ,					attention_dropout=args.attention_dropout   ,					activation_dropout=args.activation_dropout   ,					activation_function="""relu"""   ,					num_conv_layers=len(lowerCAmelCase__   )   ,					conv_channels=args.conv_channels   ,					conv_kernel_sizes=lowerCAmelCase__   ,					input_feat_per_channel=args.input_feat_per_channel   ,					input_channels=args.input_channels   ,					tie_word_embeddings=lowerCAmelCase__   ,					num_beams=5   ,					max_length=2_0_0   ,					use_cache=lowerCAmelCase__   ,					decoder_start_token_id=2   ,					early_stopping=lowerCAmelCase__   ,					)

		A__  				=       SpeechaTextForConditionalGeneration(lowerCAmelCase__   )
		A__				,      A__  				=       model.model.load_state_dict(lowerCAmelCase__   ,					strict=lowerCAmelCase__   )
		if len(lowerCAmelCase__   ) > 0 and not set(lowerCAmelCase__   ) <= {
		    "encoder.embed_positions.weights",
		    "decoder.embed_positions.weights",
		}:
				raise ValueError(
				    """Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights`  are allowed to be missing,"""
				    F' but all the following weights are missing {missing}'   )

		if tie_embeds:
				A__  				=       make_linear_from_emb(model.model.decoder.embed_tokens   )
		else:
				A__  				=       lm_head_weights

		model.save_pretrained(lowerCAmelCase__   )


if __name__ == "__main__":
					A   : List[Any]       =				argparse.ArgumentParser()
					# Required parameters
					parser.add_argument('''--fairseq_path''', type=str, help='''Path to the fairseq model (.pt) file.''')
					parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
					A   : Optional[Any]       =				parser.parse_args()
					convert_fairseq_sat_checkpoint_to_tfms(args.fairseq_path, args.pytorch_dump_folder_path)










# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from accelerate import PartialState
from accelerate.utils.operations import broadcast, gather, gather_object, pad_across_processes, reduce







def        __lowerCamelCase (							__a       :List[str]   )		->    Tuple:






		"""simple docstring"""

		return (torch.arange(state.num_processes   ) + 1.0 + (state.num_processes * state.process_index)).to(state.device   )







def        __lowerCamelCase (							__a       :int   )		->    Optional[int]:






		"""simple docstring"""

		A__  				=       create_tensor(__a   )
		A__  				=       gather(__a   )
		assert gathered_tensor.tolist() == list(range(1   ,					state.num_processes**2 + 1   )   )







def        __lowerCamelCase (							__a       :Any   )		->    Any:






		"""simple docstring"""

		A__  				=       [state.process_index]
		A__  				=       gather_object(__a   )
		assert len(__a   ) == state.num_processes, F'{gathered_obj}, {len(__a   )} != {state.num_processes}'
		assert gathered_obj == list(range(state.num_processes   )   ), F'{gathered_obj} != {list(range(state.num_processes   )   )}'







def        __lowerCamelCase (							__a       :List[str]   )		->    List[str]:






		"""simple docstring"""

		A__  				=       create_tensor(__a   )
		A__  				=       broadcast(__a   )
		assert broadcasted_tensor.shape == torch.Size([state.num_processes]   )
		assert broadcasted_tensor.tolist() == list(range(1   ,					state.num_processes + 1   )   )







def        __lowerCamelCase (							__a       :Any   )		->    Any:






		"""simple docstring"""

		if state.is_main_process:
				A__  				=       torch.arange(state.num_processes + 1   ).to(state.device   )
		else:
				A__  				=       torch.arange(state.num_processes   ).to(state.device   )
		A__  				=       pad_across_processes(__a   )
		assert padded_tensor.shape == torch.Size([state.num_processes + 1]   )
		if not state.is_main_process:
				assert padded_tensor.tolist() == list(range(0   ,					state.num_processes   )   ) + [0]







def        __lowerCamelCase (							__a       :Union[str, Any]   )		->    str:






		"""simple docstring"""

		if state.num_processes != 2:
				return
		A__  				=       create_tensor(__a   )
		A__  				=       reduce(__a   ,					"""sum"""   )
		A__  				=       torch.tensor([4.0, 6]   ).to(state.device   )
		assert torch.allclose(__a   ,					__a   ), F'{reduced_tensor} != {truth_tensor}'







def        __lowerCamelCase (							__a       :List[Any]   )		->    List[str]:






		"""simple docstring"""

		if state.num_processes != 2:
				return
		A__  				=       create_tensor(__a   )
		A__  				=       reduce(__a   ,					"""mean"""   )
		A__  				=       torch.tensor([2.0, 3]   ).to(state.device   )
		assert torch.allclose(__a   ,					__a   ), F'{reduced_tensor} != {truth_tensor}'







def        __lowerCamelCase (							__a       :str   )		->    Optional[int]:






		"""simple docstring"""

		main()







def        __lowerCamelCase (							)		->    str:






		"""simple docstring"""

		A__  				=       PartialState()
		state.print(F'State: {state}'   )
		state.print("""testing gather"""   )
		test_gather(__a   )
		state.print("""testing gather_object"""   )
		test_gather_object(__a   )
		state.print("""testing broadcast"""   )
		test_broadcast(__a   )
		state.print("""testing pad_across_processes"""   )
		test_pad_across_processes(__a   )
		state.print("""testing reduce_sum"""   )
		test_reduce_sum(__a   )
		state.print("""testing reduce_mean"""   )
		test_reduce_mean(__a   )


if __name__ == "__main__":
					main()



