






'''simple docstring'''

import json
import os
from collections import Counter

import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch import nn
from torch.utils.data import Dataset


UpperCamelCase							={1: (1, 1), 2: (2, 1), 3: (3, 1), 4: (2, 2), 5: (5, 1), 6: (3, 2), 7: (7, 1), 8: (4, 2), 9: (3, 3)}




class     A    (      nn.Module						):






 """simple docstring"""





 def __init__(							self			,     __lowerCAmelCase ):
       super().__init__()
       UpperCamelCase_				:   List[str]					    =							torchvision.models.resnetaaa(pretrained=__lowerCAmelCase )
       UpperCamelCase_				:   Union[str, Any]					    =							list(model.children() )[:-2]
       UpperCamelCase_				:   Union[str, Any]					    =							nn.Sequential(*__lowerCAmelCase )
       UpperCamelCase_				:   Optional[Any]					    =							nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds] )





 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase ):
       # Bx3x224x224 -> Bx2048x7x7 -> Bx2048xN -> BxNx2048
       UpperCamelCase_				:   int					    =							self.pool(self.model(__lowerCAmelCase ) )
       UpperCamelCase_				:   Optional[int]					    =							torch.flatten(__lowerCAmelCase			,     start_dim=2 )
       UpperCamelCase_				:   Optional[Any]					    =							out.transpose(1			,     2 ).contiguous()
       return out  # BxNx2048






class     A    (      SCREAMING_SNAKE_CASE__						):






 """simple docstring"""





 def __init__(							self			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase ):
       UpperCamelCase_				:   Dict					    =							[json.loads(__lowerCAmelCase ) for l in open(__lowerCAmelCase )]
       UpperCamelCase_				:   Union[str, Any]					    =							os.path.dirname(__lowerCAmelCase )
       UpperCamelCase_				:   Any					    =							tokenizer
       UpperCamelCase_				:   int					    =							labels
       UpperCamelCase_				:   List[str]					    =							len(__lowerCAmelCase )
       UpperCamelCase_				:   Dict					    =							max_seq_length

       UpperCamelCase_				:   int					    =							transforms
 def __len__(							self ):
       return len(self.data )
 def __getitem__(							self			,     __lowerCAmelCase ):
       UpperCamelCase_				:   Union[str, Any]					    =							torch.LongTensor(self.tokenizer.encode(self.data[index]["""text"""]			,     add_special_tokens=__lowerCAmelCase ) )
       UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_				:   Tuple					    =							sentence[0], sentence[1:-1], sentence[-1]
       UpperCamelCase_				:   List[Any]					    =							sentence[: self.max_seq_length]

       UpperCamelCase_				:   int					    =							torch.zeros(self.n_classes )
       UpperCamelCase_				:   Optional[int]					    =							1

       UpperCamelCase_				:   str					    =							Image.open(os.path.join(self.data_dir			,     self.data[index]["""img"""] ) ).convert("""RGB""" )
       UpperCamelCase_				:   List[str]					    =							self.transforms(__lowerCAmelCase )

       return {
           "image_start_token": start_token,
           "image_end_token": end_token,
           "sentence": sentence,
           "image": image,
           "label": label,
       }





 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Tuple					    =							Counter()
       for row in self.data:
             label_freqs.update(row["""label"""] )
       return label_freqs






def 							snake_case					(     a_							: Union[str, Any]	)       ->				List[Any]:





      """simple docstring"""

      UpperCamelCase_				:   List[Any]					    =							[len(row["""sentence"""]	) for row in batch]
      UpperCamelCase_    ,			UpperCamelCase_				:   Union[str, Any]					    =							len(a_	), max(a_	)

      UpperCamelCase_				:   Any					    =							torch.zeros(a_    ,   a_    ,   dtype=torch.long	)
      UpperCamelCase_				:   Optional[int]					    =							torch.zeros(a_    ,   a_    ,   dtype=torch.long	)

      for i_batch, (input_row, length) in enumerate(zip(a_    ,   a_	)	):
            UpperCamelCase_				:   Optional[int]					    =							input_row["""sentence"""]
            UpperCamelCase_				:   str					    =							1

      UpperCamelCase_				:   Dict					    =							torch.stack([row["""image"""] for row in batch]	)
      UpperCamelCase_				:   Dict					    =							torch.stack([row["""label"""] for row in batch]	)
      UpperCamelCase_				:   List[Any]					    =							torch.stack([row["""image_start_token"""] for row in batch]	)
      UpperCamelCase_				:   int					    =							torch.stack([row["""image_end_token"""] for row in batch]	)

      return text_tensor, mask_tensor, img_tensor, img_start_token, img_end_token, tgt_tensor



def 							snake_case					(     )       ->				Tuple:





      """simple docstring"""

      return [
          "Crime",
          "Drama",
          "Thriller",
          "Action",
          "Comedy",
          "Romance",
          "Documentary",
          "Short",
          "Mystery",
          "History",
          "Family",
          "Adventure",
          "Fantasy",
          "Sci-Fi",
          "Western",
          "Horror",
          "Sport",
          "War",
          "Music",
          "Musical",
          "Animation",
          "Biography",
          "Film-Noir",
      ]



def 							snake_case					(     )       ->				Any:





      """simple docstring"""

      return transforms.Compose(
          [
              transforms.Resize(256	),
              transforms.CenterCrop(224	),
              transforms.ToTensor(),
              transforms.Normalize(
                  mean=[0.46_777_044, 0.44_531_429, 0.40_661_017]    ,   std=[0.12_221_994, 0.12_145_835, 0.14_380_469]    ,   ),
          ]	)














'''simple docstring'''

import argparse
import json
import os

import fairseq
import torch
from torch import nn

from transformers import (
    SpeechaTextaConfig,
    SpeechaTextaForCausalLM,
    SpeechaTextaTokenizer,
    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 							snake_case					(     a_							: str    ,   a_							: Any    ,   a_							: Dict    ,   a_							: Optional[int]    ,   a_							: Optional[int]	)       ->				int:





      """simple docstring"""

      for attribute in key.split("""."""	):
            UpperCamelCase_				:   Union[str, Any]					    =							getattr(a_    ,   a_	)

      if weight_type is not None:
            UpperCamelCase_				:   Dict					    =							getattr(a_    ,   a_	).shape
      else:
            UpperCamelCase_				:   int					    =							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":
            UpperCamelCase_				:   str					    =							value
      elif weight_type == "weight_g":
            UpperCamelCase_				:   str					    =							value
      elif weight_type == "weight_v":
            UpperCamelCase_				:   Optional[int]					    =							value
      elif weight_type == "bias":
            UpperCamelCase_				:   Any					    =							value
      else:
            UpperCamelCase_				:   Dict					    =							value

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



def 							snake_case					(     a_							: Optional[int]    ,   a_							: int	)       ->				int:





      """simple docstring"""

      UpperCamelCase_				:   int					    =							[]
      UpperCamelCase_				:   Dict					    =							fairseq_model.state_dict()

      UpperCamelCase_				:   Union[str, Any]					    =							hf_model.feature_extractor

      # if encoder has different dim to decoder -> use proj_weight
      UpperCamelCase_				:   Optional[int]					    =							None

      for name, value in fairseq_dict.items():
            UpperCamelCase_				:   Optional[int]					    =							False
            if "conv_layers" in name:
                  load_conv_layer(
                      a_    ,   a_    ,   a_    ,   a_    ,   hf_model.config.feat_extract_norm == """group"""    ,   )
                  UpperCamelCase_				:   Any					    =							True
            elif name.split("""."""	)[0] == "proj":
                  UpperCamelCase_				:   Tuple					    =							fairseq_model.proj
                  UpperCamelCase_				:   Any					    =							True
            else:
                  for key, mapped_key in MAPPING.items():
                        if key in name or key.split("""w2v_model."""	)[-1] == name.split("""."""	)[0]:
                              UpperCamelCase_				:   Dict					    =							True
                              if "*" in mapped_key:
                                    UpperCamelCase_				:   Optional[int]					    =							name.split(a_	)[0].split("""."""	)[-2]
                                    UpperCamelCase_				:   Any					    =							mapped_key.replace("""*"""    ,   a_	)
                              if "weight_g" in name:
                                    UpperCamelCase_				:   Tuple					    =							"""weight_g"""
                              elif "weight_v" in name:
                                    UpperCamelCase_				:   Union[str, Any]					    =							"""weight_v"""
                              elif "bias" in name:
                                    UpperCamelCase_				:   int					    =							"""bias"""
                              elif "weight" in name:
                                    UpperCamelCase_				:   Optional[int]					    =							"""weight"""
                              else:
                                    UpperCamelCase_				:   int					    =							None
                              set_recursively(a_    ,   a_    ,   a_    ,   a_    ,   a_	)
                        continue
            if not is_used:
                  unused_weights.append(a_	)

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

      return proj_weight



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





      """simple docstring"""

      UpperCamelCase_				:   Optional[int]					    =							full_name.split("""conv_layers."""	)[-1]
      UpperCamelCase_				:   Optional[Any]					    =							name.split("""."""	)
      UpperCamelCase_				:   List[Any]					    =							int(items[0]	)
      UpperCamelCase_				:   Optional[Any]					    =							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."
                  )
                  UpperCamelCase_				:   List[Any]					    =							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."
                  )
                  UpperCamelCase_				:   Dict					    =							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."
                  )
                  UpperCamelCase_				:   List[str]					    =							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."
                  )
                  UpperCamelCase_				:   str					    =							value
                  logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."	)
      else:
            unused_weights.append(a_	)



def 							snake_case					(     a_							: str	)       ->				Union[str, Any]:





      """simple docstring"""

      UpperCamelCase_    ,			UpperCamelCase_				:   List[str]					    =							emb.weight.shape
      UpperCamelCase_				:   str					    =							nn.Linear(a_    ,   a_    ,   bias=a_	)
      UpperCamelCase_				:   Optional[int]					    =							emb.weight.data
      return lin_layer



def 							snake_case					(     a_							: Any	)       ->				Tuple:





      """simple docstring"""

      with open(a_    ,   """r"""    ,   encoding="""utf-8"""	) as f:
            UpperCamelCase_				:   Optional[int]					    =							f.readlines()
            UpperCamelCase_				:   Union[str, Any]					    =							[line.split(""" """	)[0] for line in lines]

      UpperCamelCase_				:   List[Any]					    =							len(a_	)

      UpperCamelCase_				:   Union[str, Any]					    =							{
          """<s>""": 0,
          """<pad>""": 1,
          """</s>""": 2,
          """<unk>""": 3,
      }

      vocab_dict.update(dict(zip(a_    ,   range(4    ,   num_words + 4	)	)	)	)
      return vocab_dict



@torch.no_grad()
def 							snake_case					(     a_							: List[Any]    ,   a_							: str    ,   a_							: int    ,   a_							: int    ,   a_							: Optional[int]    ,   a_							: Any    ,   a_							: Tuple    ,   )       ->				Optional[Any]:





      """simple docstring"""

      UpperCamelCase_				:   Dict					    =							WavaVecaConfig.from_pretrained(a_	)
      UpperCamelCase_				:   Any					    =							SpeechaTextaConfig.from_pretrained(
          a_    ,   vocab_size=a_    ,   decoder_layers=a_    ,   do_stable_layer_norm=a_	)

      UpperCamelCase_				:   Optional[Any]					    =							WavaVecaFeatureExtractor(
          feature_size=1    ,   sampling_rate=16_000    ,   padding_value=0    ,   do_normalize=a_    ,   return_attention_mask=a_    ,   )

      UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_				:   Dict					    =							fairseq.checkpoint_utils.load_model_ensemble_and_task(
          [checkpoint_path]    ,   arg_overrides={"""data""": """/""".join(dict_path.split("""/"""	)[:-1]	)}	)
      UpperCamelCase_				:   Union[str, Any]					    =							model[0].eval()

      # set weights for wav2vec2 encoder
      UpperCamelCase_				:   Tuple					    =							WavaVecaModel(a_	)
      UpperCamelCase_				:   Any					    =							recursively_load_weights_wavaveca(model.encoder    ,   a_	)

      UpperCamelCase_				:   Any					    =							SpeechaTextaForCausalLM(a_	)
      UpperCamelCase_    ,			UpperCamelCase_				:   Optional[Any]					    =							hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()    ,   strict=a_	)

      # set output linear layer
      unexpected_keys.remove("""embed_out"""	)
      UpperCamelCase_				:   Tuple					    =							nn.Parameter(model.decoder.embed_out.detach()	)

      # layer norm is init to identity matrix so leaving it is fine
      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}"	)

      UpperCamelCase_				:   Optional[Any]					    =							SpeechEncoderDecoderModel(encoder=a_    ,   decoder=a_	)
      UpperCamelCase_				:   Dict					    =							False

      # add projection layer
      UpperCamelCase_				:   Any					    =							nn.Parameter(projection_layer.weight	)
      UpperCamelCase_				:   Optional[Any]					    =							nn.Parameter(projection_layer.bias	)

      UpperCamelCase_				:   Dict					    =							create_vocab_dict(a_	)

      with open(os.path.join(a_    ,   """vocab.json"""	)    ,   """w"""	) as fp:
            json.dump(a_    ,   a_	)

      UpperCamelCase_				:   Optional[int]					    =							SpeechaTextaTokenizer(os.path.join(a_    ,   """vocab.json"""	)	)
      tokenizer.save_pretrained(a_	)

      UpperCamelCase_				:   Optional[int]					    =							hf_wavavec.config.to_dict()
      UpperCamelCase_				:   Union[str, Any]					    =							tokenizer.pad_token_id
      UpperCamelCase_				:   List[Any]					    =							tokenizer.bos_token_id
      UpperCamelCase_				:   str					    =							tokenizer.eos_token_id
      UpperCamelCase_				:   Dict					    =							"""speech_to_text_2"""
      UpperCamelCase_				:   Optional[Any]					    =							"""wav2vec2"""

      UpperCamelCase_				:   Dict					    =							SpeechEncoderDecoderConfig.from_dict(a_	)

      hf_wavavec.save_pretrained(a_	)
      feature_extractor.save_pretrained(a_	)


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(
          "--encoder_config_path",
          default="facebook/wav2vec2-large-lv60",
          type=str,
          help="Path to hf encoder wav2vec2 checkpoint config",
      )
      parser.add_argument(
          "--decoder_config_path",
          default="facebook/s2t-small-mustc-en-fr-st",
          type=str,
          help="Path to hf decoder s2t checkpoint config",
      )
      parser.add_argument("--vocab_size", default=1_0224, type=int, help="Vocab size of decoder")
      parser.add_argument("--num_decoder_layers", default=7, type=int, help="Number of decoder layers")

      UpperCamelCase							=parser.parse_args()
      convert_wavaveca_checkpoint(
          args.checkpoint_path,
          args.pytorch_dump_folder_path,
          args.dict_path,
          encoder_config_path=args.encoder_config_path,
          decoder_config_path=args.decoder_config_path,
          vocab_size=args.vocab_size,
          num_decoder_layers=args.num_decoder_layers,
      )






