


import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
UpperCAmelCase      =   logging.get_logger(__name__)

UpperCAmelCase      =   {
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}
UpperCAmelCase      =   [
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Any:




      """simple docstring"""

      for attribute in key.split('''.''' ):
            snake_case_	=				getattr(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

      if weight_type is not None:
            snake_case_	=				getattr(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE ).shape
      else:
            snake_case_	=				hf_pointer.shape

      assert hf_shape == value.shape, (
          f'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
          f''' {value.shape} for {full_name}'''
      )

      if weight_type == "weight":
            snake_case_	=				value
      elif weight_type == "weight_g":
            snake_case_	=				value
      elif weight_type == "weight_v":
            snake_case_	=				value
      elif weight_type == "bias":
            snake_case_	=				value
      else:
            snake_case_	=				value

      logger.info(f'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.''' )






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Optional[int]:




      """simple docstring"""

      snake_case_	=				[]
      snake_case_	=				fairseq_model.state_dict()

      snake_case_	=				hf_model.feature_extractor
      snake_case_	=				hf_model.adapter

      for name, value in fairseq_dict.items():
            snake_case_	=				False
            if "conv_layers" in name:
                  load_conv_layer(
                      SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , hf_model.config.feat_extract_norm == '''group''' , )
                  snake_case_	=				True
            elif any(x in name for x in ['''adaptor''', '''w2v_encoder.proj.''', '''w2v_proj_ln.'''] ):
                  load_adapter(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )
                  snake_case_	=				True
            else:
                  for key, mapped_key in MAPPING.items():
                        if key in name or key.split('''w2v_model.''' )[-1] == name.split('''.''' )[0]:
                              snake_case_	=				True
                              if "*" in mapped_key:
                                    snake_case_	=				name.split(SCREAMING_SNAKE_CASE )[0].split('''.''' )[-2]
                                    snake_case_	=				mapped_key.replace('''*''' , SCREAMING_SNAKE_CASE )
                              if "weight_g" in name:
                                    snake_case_	=				'''weight_g'''
                              elif "weight_v" in name:
                                    snake_case_	=				'''weight_v'''
                              elif "bias" in name:
                                    snake_case_	=				'''bias'''
                              elif "weight" in name:
                                    snake_case_	=				'''weight'''
                              else:
                                    snake_case_	=				None
                              set_recursively(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )
                        continue
            if not is_used:
                  unused_weights.append(SCREAMING_SNAKE_CASE )

      logger.warning(f'''Unused weights: {unused_weights}''' )






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  List[Any]:




      """simple docstring"""

      snake_case_	=				full_name.split('''conv_layers.''' )[-1]
      snake_case_	=				name.split('''.''' )
      snake_case_	=				int(items[0] )
      snake_case_	=				int(items[1] )

      if type_id == 0:
            if "bias" in name:
                  assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
                      f'''{full_name} has size {value.shape}, but'''
                      f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
                  )
                  snake_case_	=				value
                  logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.''' )
            elif "weight" in name:
                  assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
                      f'''{full_name} has size {value.shape}, but'''
                      f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
                  )
                  snake_case_	=				value
                  logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.''' )
      elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
            if "bias" in name:
                  assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
                      f'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
                      " found."
                  )
                  snake_case_	=				value
                  logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
            elif "weight" in name:
                  assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
                      f'''{full_name} has size {value.shape}, but'''
                      f''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
                  )
                  snake_case_	=				value
                  logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
      else:
            unused_weights.append(SCREAMING_SNAKE_CASE )






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  str:




      """simple docstring"""

      snake_case_	=				full_name.split('''adaptor.''' )[-1]
      snake_case_	=				name.split('''.''' )

      if items[1].isdigit():
            snake_case_	=				int(items[1] )
      else:
            snake_case_	=				None

      if "adaptor" not in full_name:
            if "proj_ln" in full_name:
                  # has to be layer norm
                  if "bias" in name:
                        assert (
                            value.shape == adapter.proj_layer_norm.bias.data.shape
                        ), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
                        snake_case_	=				value
                        logger.info(f'''Adapter proj layer norm bias was initialized from {full_name}.''' )
                  if "weight" in name:
                        assert (
                            value.shape == adapter.proj_layer_norm.weight.data.shape
                        ), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
                        snake_case_	=				value
            else:
                  # has to be projection layer
                  if "bias" in name:
                        assert (
                            value.shape == adapter.proj.bias.data.shape
                        ), f'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
                        snake_case_	=				value
                        logger.info(f'''Adapter proj layer bias was initialized from {full_name}.''' )
                  if "weight" in name:
                        assert (
                            value.shape == adapter.proj.weight.data.shape
                        ), f'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
                        snake_case_	=				value
                        logger.info(f'''Adapter proj layer weight was initialized from {full_name}.''' )
      elif isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE ):
            if "bias" in name:
                  assert (
                      value.shape == adapter.layers[layer_id].conv.bias.data.shape
                  ), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
                  snake_case_	=				value
                  logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
            elif "weight" in name:
                  assert (
                      value.shape == adapter.layers[layer_id].conv.weight.data.shape
                  ), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
                  snake_case_	=				value
                  logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
      else:
            unused_weights.append(SCREAMING_SNAKE_CASE )






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  Any:




      """simple docstring"""

      snake_case_	,   snake_case_	=				emb.weight.shape
      snake_case_	=				nn.Linear(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , bias=SCREAMING_SNAKE_CASE )
      snake_case_	=				emb.weight.data
      return lin_layer






@torch.no_grad()
def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , )->  Optional[int]:




      """simple docstring"""

      snake_case_	=				WavaVecaConfig.from_pretrained(
          SCREAMING_SNAKE_CASE , add_adapter=SCREAMING_SNAKE_CASE , adapter_stride=SCREAMING_SNAKE_CASE , adapter_kernel_size=SCREAMING_SNAKE_CASE , use_auth_token=SCREAMING_SNAKE_CASE , output_hidden_size=SCREAMING_SNAKE_CASE , )
      snake_case_	=				MBartConfig.from_pretrained(SCREAMING_SNAKE_CASE )

      # load model
      snake_case_	,   snake_case_	,   snake_case_	=				fairseq.checkpoint_utils.load_model_ensemble_and_task(
          [checkpoint_path] , arg_overrides={
              '''config_yaml''': config_yaml_path,
              '''data''': '''/'''.join(dict_path.split('''/''' )[:-1] ),
              '''w2v_path''': checkpoint_path,
              '''load_pretrained_decoder_from''': None,
          } , )
      snake_case_	=				model[0].eval()

      # load feature extractor
      snake_case_	=				WavaVecaFeatureExtractor.from_pretrained(SCREAMING_SNAKE_CASE , use_auth_token=SCREAMING_SNAKE_CASE )

      # set weights for wav2vec2 encoder
      snake_case_	=				WavaVecaModel(SCREAMING_SNAKE_CASE )

      recursively_load_weights_wavaveca(model.encoder , SCREAMING_SNAKE_CASE )

      # load decoder weights
      snake_case_	=				MBartForCausalLM(SCREAMING_SNAKE_CASE )
      snake_case_	,   snake_case_	=				hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() , strict=SCREAMING_SNAKE_CASE )
      logger.warning(f'''The following keys are missing when loading the decoder weights: {missing_keys}''' )
      logger.warning(f'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}''' )

      snake_case_	=				SpeechEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE , decoder=SCREAMING_SNAKE_CASE )
      snake_case_	=				False

      snake_case_	=				MBartaaTokenizer(SCREAMING_SNAKE_CASE )
      tokenizer.save_pretrained(SCREAMING_SNAKE_CASE )

      snake_case_	=				hf_wavavec.config.to_dict()
      snake_case_	=				tokenizer.pad_token_id
      snake_case_	=				tokenizer.bos_token_id
      snake_case_	=				tokenizer.eos_token_id
      snake_case_	=				'''mbart50'''
      snake_case_	=				'''wav2vec2'''

      snake_case_	=				tokenizer.eos_token_id
      snake_case_	=				25_0004
      snake_case_	=				tokenizer.eos_token_id

      snake_case_	=				SpeechEncoderDecoderConfig.from_dict(SCREAMING_SNAKE_CASE )

      hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE )
      feature_extractor.save_pretrained(SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
   UpperCAmelCase      =   argparse.ArgumentParser()
   parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
   parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
   parser.add_argument("""--dict_path""", default=None, type=str, help="""Path to dict of fine-tuned model""")
   parser.add_argument("""--config_yaml_path""", default=None, type=str, help="""Path to yaml file of fine-tuned model""")
   parser.add_argument(
       """--encoder_config_path""",
       default="""facebook/wav2vec2-xls-r-1b""",
       type=str,
       help="""Path to hf encoder wav2vec2 checkpoint config""",
   )
   parser.add_argument(
       """--decoder_config_path""",
       default="""facebook/mbart-large-50-one-to-many-mmt""",
       type=str,
       help="""Path to hf decoder checkpoint config""",
   )
   parser.add_argument("""--add_adapter""", default=True, type=bool, help="""whethere to add model adapter layers""")
   parser.add_argument("""--adapter_stride""", default=2, type=int, help="""stride of adapter layers""")
   parser.add_argument("""--adapter_kernel_size""", default=3, type=int, help="""kernel size of adapter layers""")
   parser.add_argument("""--encoder_output_dim""", default=1024, type=int, help="""encoder output dim""")
   parser.add_argument("""--start_token_id""", default=25_0004, type=int, help="""`decoder_start_token_id` of model config""")

   UpperCAmelCase      =   parser.parse_args()
   convert_wavaveca_checkpoint(
       args.checkpoint_path,
       args.pytorch_dump_folder_path,
       args.dict_path,
       args.config_yaml_path,
       encoder_config_path=args.encoder_config_path,
       decoder_config_path=args.decoder_config_path,
       add_adapter=args.add_adapter,
       adapter_kernel_size=args.adapter_kernel_size,
       adapter_stride=args.adapter_stride,
       decoder_start_token_id=args.start_token_id,
       encoder_output_dim=args.encoder_output_dim,
   )



from __future__ import annotations

import time

UpperCAmelCase      =   list[tuple[int, int]]

UpperCAmelCase      =   [
    [0, 0, 0, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
    [0, 0, 0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0, 0, 0],
    [1, 0, 1, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0],
]

UpperCAmelCase      =   [[-1, 0], [0, -1], [1, 0], [0, 1]]  # up, left, down, right







class lowerCAmelCase_						:

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				pos_x
       snake_case_	=				pos_y
       snake_case_	=				(pos_y, pos_x)
       snake_case_	=				goal_x
       snake_case_	=				goal_y
       snake_case_	=				parent







class lowerCAmelCase_						:

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				Node(start[1]		,				start[0]		,				goal[1]		,				goal[0]		,				_UpperCAmelCase )
       snake_case_	=				Node(goal[1]		,				goal[0]		,				goal[1]		,				goal[0]		,				_UpperCAmelCase )

       snake_case_	=				[self.start]
       snake_case_	=				False



 def 						UpperCamelCase__					(     self ):
       while self.node_queue:
             snake_case_	=				self.node_queue.pop(0 )

             if current_node.pos == self.target.pos:
                   snake_case_	=				True
                   return self.retrace_path(_UpperCAmelCase )

             snake_case_	=				self.get_successors(_UpperCAmelCase )

             for node in successors:
                   self.node_queue.append(_UpperCAmelCase )

       if not self.reached:
             return [self.start.pos]
       return None



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       snake_case_	=				[]
       for action in delta:
             snake_case_	=				parent.pos_x + action[1]
             snake_case_	=				parent.pos_y + action[0]
             if not (0 <= pos_x <= len(grid[0] ) - 1 and 0 <= pos_y <= len(_UpperCAmelCase ) - 1):
                   continue

             if grid[pos_y][pos_x] != 0:
                   continue

             successors.append(
                 Node(_UpperCAmelCase		,				_UpperCAmelCase		,				self.target.pos_y		,				self.target.pos_x		,				_UpperCAmelCase ) )
       return successors





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       snake_case_	=				node
       snake_case_	=				[]
       while current_node is not None:
             path.append((current_node.pos_y, current_node.pos_x) )
             snake_case_	=				current_node.parent
       path.reverse()
       return path







class lowerCAmelCase_						:

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				BreadthFirstSearch(_UpperCAmelCase		,				_UpperCAmelCase )
       snake_case_	=				BreadthFirstSearch(_UpperCAmelCase		,				_UpperCAmelCase )
       snake_case_	=				False



 def 						UpperCamelCase__					(     self ):
       while self.fwd_bfs.node_queue or self.bwd_bfs.node_queue:
             snake_case_	=				self.fwd_bfs.node_queue.pop(0 )
             snake_case_	=				self.bwd_bfs.node_queue.pop(0 )

             if current_bwd_node.pos == current_fwd_node.pos:
                   snake_case_	=				True
                   return self.retrace_bidirectional_path(
                       _UpperCAmelCase		,				_UpperCAmelCase )

             snake_case_	=				current_bwd_node
             snake_case_	=				current_fwd_node

             snake_case_	=				{
                 self.fwd_bfs: self.fwd_bfs.get_successors(_UpperCAmelCase ),
                 self.bwd_bfs: self.bwd_bfs.get_successors(_UpperCAmelCase ),
             }

             for bfs in [self.fwd_bfs, self.bwd_bfs]:
                   for node in successors[bfs]:
                         bfs.node_queue.append(_UpperCAmelCase )

       if not self.reached:
             return [self.fwd_bfs.start.pos]
       return None





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				self.fwd_bfs.retrace_path(_UpperCAmelCase )
       snake_case_	=				self.bwd_bfs.retrace_path(_UpperCAmelCase )
       bwd_path.pop()
       bwd_path.reverse()
       snake_case_	=				fwd_path + bwd_path
       return path


if __name__ == "__main__":
   # all coordinates are given in format [y,x]
   import doctest

   doctest.testmod()
   UpperCAmelCase      =   (0, 0)
   UpperCAmelCase      =   (len(grid) - 1, len(grid[0]) - 1)
   for elem in grid:
      print(elem)

   UpperCAmelCase      =   time.time()
   UpperCAmelCase      =   BreadthFirstSearch(init, goal)
   UpperCAmelCase      =   bfs.search()
   UpperCAmelCase      =   time.time() - start_bfs_time

   print("""Unidirectional BFS computation time : """, bfs_time)

   UpperCAmelCase      =   time.time()
   UpperCAmelCase      =   BidirectionalBreadthFirstSearch(init, goal)
   UpperCAmelCase      =   bd_bfs.search()
   UpperCAmelCase      =   time.time() - start_bd_bfs_time

   print("""Bidirectional BFS computation time : """, bd_bfs_time)