



'''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()
_lowerCamelCase	:       int      	=	logging.get_logger(__name__)

_lowerCamelCase	:       Optional[Any]      	=	{
    "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",
}
_lowerCamelCase	:       Union[str, Any]      	=	[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]







def 				__lowerCamelCase				(   A__    ,						A__    ,						A__    ,						A__    ,						A__    )       ->	Any:






   """simple docstring"""
   for attribute in key.split('.'    ):
      UpperCamelCase					=  getattr(A__    ,						A__    )

   if weight_type is not None:
      UpperCamelCase					=  getattr(A__    ,						A__    ).shape
   else:
      UpperCamelCase					=  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					=  value
   elif weight_type == "weight_g":
      UpperCamelCase					=  value
   elif weight_type == "weight_v":
      UpperCamelCase					=  value
   elif weight_type == "bias":
      UpperCamelCase					=  value
   else:
      UpperCamelCase					=  value

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







def 				__lowerCamelCase				(   A__    ,						A__    )       ->	Dict:






   """simple docstring"""
   UpperCamelCase					=  []
   UpperCamelCase					=  fairseq_model.state_dict()

   UpperCamelCase					=  hf_model.feature_extractor

   # if encoder has different dim to decoder -> use proj_weight
   UpperCamelCase					=  None

   for name, value in fairseq_dict.items():
      UpperCamelCase					=  False
      if "conv_layers" in name:
         load_conv_layer(
             A__    ,						A__    ,						A__    ,						A__    ,						hf_model.config.feat_extract_norm == 'group'    ,						)
         UpperCamelCase					=  True
      elif name.split('.'    )[0] == "proj":
         UpperCamelCase					=  fairseq_model.proj
         UpperCamelCase					=  True
      else:
         for key, mapped_key in MAPPING.items():
            if key in name or key.split('w2v_model.'    )[-1] == name.split('.'    )[0]:
               UpperCamelCase					=  True
               if "*" in mapped_key:
                  UpperCamelCase					=  name.split(A__    )[0].split('.'    )[-2]
                  UpperCamelCase					=  mapped_key.replace('*'    ,						A__    )
               if "weight_g" in name:
                  UpperCamelCase					=  'weight_g'
               elif "weight_v" in name:
                  UpperCamelCase					=  'weight_v'
               elif "bias" in name:
                  UpperCamelCase					=  'bias'
               elif "weight" in name:
                  UpperCamelCase					=  'weight'
               else:
                  UpperCamelCase					=  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 				__lowerCamelCase				(   A__    ,						A__    ,						A__    ,						A__    ,						A__    )       ->	str:






   """simple docstring"""
   UpperCamelCase					=  full_name.split('conv_layers.'    )[-1]
   UpperCamelCase					=  name.split('.'    )
   UpperCamelCase					=  int(items[0]    )
   UpperCamelCase					=  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					=  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					=  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					=  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					=  value
         logger.info(F"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""    )
   else:
      unused_weights.append(A__    )







def 				__lowerCamelCase				(   A__    )       ->	int:






   """simple docstring"""
   UpperCamelCase				,      UpperCamelCase					=  emb.weight.shape
   UpperCamelCase					=  nn.Linear(A__    ,						A__    ,						bias=A__    )
   UpperCamelCase					=  emb.weight.data
   return lin_layer







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






   """simple docstring"""
   with open(A__    ,						'r'    ,						encoding='utf-8'    ) as f:
      UpperCamelCase					=  f.readlines()
      UpperCamelCase					=  [line.split(' '    )[0] for line in lines]

   UpperCamelCase					=  len(A__    )

   UpperCamelCase					=  {
       '<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 				__lowerCamelCase				(   A__    ,						A__    ,						A__    ,						A__    ,						A__    ,						A__    ,						A__    ,						)       ->	Optional[int]:






   """simple docstring"""
   UpperCamelCase					=  WavaVecaConfig.from_pretrained(A__    )
   UpperCamelCase					=  SpeechaTextaConfig.from_pretrained(
       A__    ,						vocab_size=A__    ,						decoder_layers=A__    ,						do_stable_layer_norm=A__    )

   UpperCamelCase					=  WavaVecaFeatureExtractor(
       feature_size=1    ,						sampling_rate=16_000    ,						padding_value=0    ,						do_normalize=A__    ,						return_attention_mask=A__    ,						)

   UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  fairseq.checkpoint_utils.load_model_ensemble_and_task(
       [checkpoint_path]    ,						arg_overrides={'data': '/'.join(dict_path.split('/'    )[:-1]    )}    )
   UpperCamelCase					=  model[0].eval()

   # set weights for wav2vec2 encoder
   UpperCamelCase					=  WavaVecaModel(A__    )
   UpperCamelCase					=  recursively_load_weights_wavaveca(model.encoder    ,						A__    )

   UpperCamelCase					=  SpeechaTextaForCausalLM(A__    )
   UpperCamelCase				,      UpperCamelCase					=  hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()    ,						strict=A__    )

   # set output linear layer
   unexpected_keys.remove('embed_out'    )
   UpperCamelCase					=  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					=  SpeechEncoderDecoderModel(encoder=A__    ,						decoder=A__    )
   UpperCamelCase					=  False

   # add projection layer
   UpperCamelCase					=  nn.Parameter(projection_layer.weight    )
   UpperCamelCase					=  nn.Parameter(projection_layer.bias    )

   UpperCamelCase					=  create_vocab_dict(A__    )

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

   UpperCamelCase					=  SpeechaTextaTokenizer(os.path.join(A__    ,						'vocab.json'    )    )
   tokenizer.save_pretrained(A__    )

   UpperCamelCase					=  hf_wavavec.config.to_dict()
   UpperCamelCase					=  tokenizer.pad_token_id
   UpperCamelCase					=  tokenizer.bos_token_id
   UpperCamelCase					=  tokenizer.eos_token_id
   UpperCamelCase					=  'speech_to_text_2'
   UpperCamelCase					=  'wav2vec2'

   UpperCamelCase					=  SpeechEncoderDecoderConfig.from_dict(A__    )

   hf_wavavec.save_pretrained(A__    )
   feature_extractor.save_pretrained(A__    )


if __name__ == "__main__":
   _lowerCamelCase	:       Dict      	=	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")

   _lowerCamelCase	:       List[Any]      	=	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,
   )






'''simple docstring'''


from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule







def 				__lowerCamelCase				(   A__    ,						A__    ,						A__=1e-1_2    )       ->	Dict:






   """simple docstring"""
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   return jnp.matmul(A__    ,						norm_emb_a.T    )

class  SCREAMING_SNAKE_CASE (  nn.Module		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  42
    _SCREAMING_SNAKE_CASE	  =  jnp.floataa






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  FlaxCLIPVisionModule(self.config.vision_config )
       UpperCamelCase					=  nn.Dense(self.config.projection_dim  ,						use_bias=UpperCamelCase__  ,						dtype=self.dtype )

       UpperCamelCase					=  self.param('concept_embeds'  ,						jax.nn.initializers.ones  ,						(1_7, self.config.projection_dim) )
       UpperCamelCase					=  self.param(
           'special_care_embeds'  ,						jax.nn.initializers.ones  ,						(3, self.config.projection_dim) )

       UpperCamelCase					=  self.param('concept_embeds_weights'  ,						jax.nn.initializers.ones  ,						(1_7,) )
       UpperCamelCase					=  self.param('special_care_embeds_weights'  ,						jax.nn.initializers.ones  ,						(3,) )






    def __call__(  self      :					str  ,						UpperCamelCase__      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  self.vision_model(UpperCamelCase__ )[1]
       UpperCamelCase					=  self.visual_projection(UpperCamelCase__ )

       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.special_care_embeds )
       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.concept_embeds )

       # increase this value to create a stronger `nfsw` filter
       # at the cost of increasing the possibility of filtering benign image inputs
       UpperCamelCase					=  0.0

       UpperCamelCase					=  special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
       UpperCamelCase					=  jnp.round(UpperCamelCase__  ,						3 )
       UpperCamelCase					=  jnp.any(special_scores > 0  ,						axis=1  ,						keepdims=UpperCamelCase__ )
       # Use a lower threshold if an image has any special care concept
       UpperCamelCase					=  is_special_care * 0.0_1

       UpperCamelCase					=  cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
       UpperCamelCase					=  jnp.round(UpperCamelCase__  ,						3 )
       UpperCamelCase					=  jnp.any(concept_scores > 0  ,						axis=1 )

       return has_nsfw_concepts

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  CLIPConfig
    _SCREAMING_SNAKE_CASE	  =  """clip_input"""
    _SCREAMING_SNAKE_CASE	  =  FlaxStableDiffusionSafetyCheckerModule






    def __init__(  self      :					Union[str, Any]  ,						UpperCamelCase__      :					CLIPConfig  ,						UpperCamelCase__      :					Optional[Tuple] = None  ,						UpperCamelCase__      :					int = 0  ,						UpperCamelCase__      :					jnp.dtype = jnp.floataa  ,						UpperCamelCase__      :					bool = True  ,						**UpperCamelCase__      :					List[str]  ,						):


       """simple docstring"""







       if input_shape is None:
          UpperCamelCase					=  (1, 2_2_4, 2_2_4, 3)
       UpperCamelCase					=  self.module_class(config=UpperCamelCase__  ,						dtype=UpperCamelCase__  ,						**UpperCamelCase__ )
       super().__init__(UpperCamelCase__  ,						UpperCamelCase__  ,						input_shape=UpperCamelCase__  ,						seed=UpperCamelCase__  ,						dtype=UpperCamelCase__  ,						_do_init=_do_init )






    def  A       (  self      :					int  ,						UpperCamelCase__      :					jax.random.KeyArray  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					FrozenDict = None ):


       """simple docstring"""







       UpperCamelCase					=  jax.random.normal(UpperCamelCase__  ,						UpperCamelCase__ )

       UpperCamelCase				,      UpperCamelCase					=  jax.random.split(UpperCamelCase__ )
       UpperCamelCase					=  {'params': params_rng, 'dropout': dropout_rng}

       UpperCamelCase					=  self.module.init(UpperCamelCase__  ,						UpperCamelCase__ )['params']

       return random_params






    def __call__(  self      :					List[Any]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					dict = None  ,						):


       """simple docstring"""







       UpperCamelCase					=  jnp.transpose(UpperCamelCase__  ,						(0, 2, 3, 1) )

       return self.module.apply(
           {'params': params or self.params}  ,						jnp.array(UpperCamelCase__  ,						dtype=jnp.floataa )  ,						rngs={}  ,						)

