



'''simple docstring'''



import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


a__						:			Dict										=			{0: 'null', 1: 'short', 2: 'long', 3: 'yes', 4: 'no'}
a__						:			List[str]										=			set(''.join(['‘', '’', '´', '`', '.', ',', '-', '"']))







def _lowercase   (  __A							,__A=0							,__A=None       ):



 '''simple docstring'''



 return [" ".join(x.split(""" """       )[begin:end]       ) for x in answers if len(x.split(""" """       )       ) > 1]







def _lowercase   (  __A							,__A=False       ):



 '''simple docstring'''



 if make_sub_answers:
  # if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
  # *e.g.* if the correct answer contains a prefix such as "the", or "a"
  __UpperCamelCase				      =  (
      given_answers + get_sub_answers(__A							,begin=1       ) + get_sub_answers(__A							,end=-1       )
  )
 __UpperCamelCase				      =  []
 for answer in given_answers:
  __UpperCamelCase				      =  answer.replace("""_"""							,""" """       ).lower()
  __UpperCamelCase				      =  """""".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else """ """ for c in alias       )
  answers.append(""" """.join(alias.split()       ).strip()       )
 return set(__A       )







def _lowercase   (  __A							,__A							,__A=1							,__A=100       ):



 '''simple docstring'''



 __UpperCamelCase      ,		__UpperCamelCase				      =  jax.lax.top_k(__A							,__A       )
 __UpperCamelCase      ,		__UpperCamelCase				      =  jax.lax.top_k(__A							,__A       )

 __UpperCamelCase				      =  best_end_idx[:, None] - best_start_idx[None, :]
 __UpperCamelCase				      =  jnp.logical_or(widths < 0							,widths > max_size       )
 __UpperCamelCase				      =  (best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
 __UpperCamelCase				      =  jnp.argmax(__A       ).item()

 return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]







def _lowercase   (  __A       ):



 '''simple docstring'''



 __UpperCamelCase				      =  sample["""question"""]["""text"""]
 __UpperCamelCase				      =  sample["""document"""]["""tokens"""]["""token"""]
 __UpperCamelCase				      =  sample["""document"""]["""tokens"""]["""is_html"""]
 __UpperCamelCase				      =  sample["""annotations"""]["""long_answer"""]
 __UpperCamelCase				      =  sample["""annotations"""]["""short_answers"""]

 __UpperCamelCase				      =  """ """.join([context[i] for i in range(len(__A       )       ) if not is_html[i]]       )

 # 0 - No ; 1 - Yes
 for answer in sample["annotations"]["yes_no_answer"]:
  if answer == 0 or answer == 1:
   return {
       "question": question,
       "context": context_string,
       "short": [],
       "long": [],
       "category": "no" if answer == 0 else "yes",
   }

 __UpperCamelCase				      =  []
 for s in short_answers:
  short_targets.extend(s["""text"""]       )
 __UpperCamelCase				      =  list(set(__A       )       )

 __UpperCamelCase				      =  []
 for s in long_answers:
  if s["start_token"] == -1:
   continue
  __UpperCamelCase				      =  context[s["""start_token"""] : s["""end_token"""]]
  __UpperCamelCase				      =  is_html[s["""start_token"""] : s["""end_token"""]]
  __UpperCamelCase				      =  """ """.join([answer[i] for i in range(len(__A       )       ) if not html[i]]       )
  if new_answer not in long_targets:
   long_targets.append(__A       )

 __UpperCamelCase				      =  """long_short""" if len(short_targets + long_targets       ) > 0 else """null"""

 return {
     "question": question,
     "context": context_string,
     "short": short_targets,
     "long": long_targets,
     "category": category,
 }







def _lowercase   (  ):



 '''simple docstring'''



 __UpperCamelCase				      =  load_from_disk("""natural-questions-validation"""       )
 __UpperCamelCase				      =  dataset.map(__A       ).remove_columns(["""annotations""", """document""", """id"""]       )
 print(__A       )

 __UpperCamelCase				      =  dataset.filter(lambda __A       : (len(x["""question"""]       ) + len(x["""context"""]       )) < 4 * 4_096       )
 __UpperCamelCase				      =  short_validation_dataset.filter(lambda __A       : x["category"] != "null"       )
 short_validation_dataset

 __UpperCamelCase				      =  """vasudevgupta/flax-bigbird-natural-questions"""
 __UpperCamelCase				      =  FlaxBigBirdForNaturalQuestions.from_pretrained(__A       )
 __UpperCamelCase				      =  BigBirdTokenizerFast.from_pretrained(__A       )

 @jax.jit
 def forward(*__A							,**__A       ):
  __UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase				      =  model(*__A							,**__A       )
  return start_logits, end_logits, jnp.argmax(__A							,axis=-1       )

 def evaluate(__A       ):
  # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
  __UpperCamelCase				      =  tokenizer(
      example["""question"""]							,example["""context"""]							,return_tensors="""np"""							,max_length=4_096							,padding="""max_length"""							,truncation=__A							,)

  __UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase				      =  forward(**__A       )

  __UpperCamelCase				      =  CATEGORY_MAPPING[category.item()]

  __UpperCamelCase				      =  example["""long"""] + example["""short"""]
  if example["category"] in ["yes", "no", "null"]:
   __UpperCamelCase				      =  [example["""category"""]]
  __UpperCamelCase				      =  example["""category"""] != """null"""
  # Now target can be: "yes", "no", "null", "list of long & short answers"

  if predicted_category in ["yes", "no", "null"]:
   __UpperCamelCase				      =  [predicted_category]
   __UpperCamelCase				      =  example["""output"""] == example["""targets"""]
   __UpperCamelCase				      =  predicted_category != """null"""
   return example

  __UpperCamelCase				      =  38 if predicted_category == """short""" else 1_024
  __UpperCamelCase      ,		__UpperCamelCase				      =  get_best_valid_start_end_idx(
      start_scores[0]							,end_scores[0]							,top_k=8							,max_size=__A       )

  __UpperCamelCase				      =  inputs["""input_ids"""][0].tolist()
  __UpperCamelCase				      =  [tokenizer.decode(input_ids[start_score : end_score + 1]       )]

  __UpperCamelCase				      =  expand_to_aliases(example["""targets"""]							,make_sub_answers=__A       )
  __UpperCamelCase				      =  expand_to_aliases(example["""output"""]       )

  # some preprocessing to both prediction and answer
  __UpperCamelCase				      =  {"""""".join(a.split()       ) for a in answers}
  __UpperCamelCase				      =  {"""""".join(p.split()       ) for p in predictions}
  __UpperCamelCase				      =  {s for s in predictions if s not in ["""``""", """''""", """`""", """'"""]}

  # if there is a common element, it's a exact match
  __UpperCamelCase				      =  len(list(answers & predictions       )       ) > 0
  __UpperCamelCase				      =  predicted_category != """null""" and len(__A       ) > 0

  return example

 __UpperCamelCase				      =  short_validation_dataset.map(__A       )

 __UpperCamelCase				      =  len(__A       )
 __UpperCamelCase				      =  len(short_validation_dataset.filter(lambda __A       : x["match"] == 1       )       )
 print("""EM score:"""							,(matched / total) * 100							,"""%"""       )


if __name__ == "__main__":
 main()










'''simple docstring'''



import argparse

import torch
from saprl.hub import distilhubert

from transformers import HubertConfig, HubertModel, WavaVecaFeatureExtractor, logging


logging.set_verbosity_info()
a__						:			str										=			logging.get_logger(__name__)

a__						:			Optional[int]										=			{
    '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',
    'mask_emb': 'masked_spec_embed',
}







def _lowercase   (  __A							,__A							,__A							,__A							,__A       ):



 '''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 _lowercase   (  __A							,__A       ):



 '''simple docstring'''



 __UpperCamelCase				      =  []
 __UpperCamelCase				      =  fairseq_model.state_dict()

 __UpperCamelCase				      =  hf_model.feature_extractor

 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
  else:
   for key, mapped_key in MAPPING.items():
    __UpperCamelCase				      =  mapped_key

    if key in name:
     __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 "weight" in name:
      __UpperCamelCase				      =  """weight"""
     elif "bias" in name:
      __UpperCamelCase				      =  """bias"""
     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}"       )







def _lowercase   (  __A							,__A							,__A							,__A							,__A       ):



 '''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 _lowercase   (  __A       ):



 '''simple docstring'''



 __UpperCamelCase				      =  HubertConfig()
 __UpperCamelCase				      =  model.config

 __UpperCamelCase				      =  fs_config.activation_dropout
 __UpperCamelCase				      =  False
 __UpperCamelCase				      =  fs_config.attention_dropout
 __UpperCamelCase				      =  False
 __UpperCamelCase				      =  eval(fs_config.extractor_conv_feature_layers       )
 __UpperCamelCase				      =  [x[0] for x in conv_layers]
 __UpperCamelCase				      =  [x[1] for x in conv_layers]
 __UpperCamelCase				      =  [x[2] for x in conv_layers]
 __UpperCamelCase				      =  """gelu"""
 __UpperCamelCase				      =  """layer""" if fs_config.extractor_mode == """layer_norm""" else """group"""
 __UpperCamelCase				      =  False
 __UpperCamelCase				      =  0.0
 __UpperCamelCase				      =  0.0
 __UpperCamelCase				      =  fs_config.activation_fn
 __UpperCamelCase				      =  fs_config.dropout
 __UpperCamelCase				      =  fs_config.encoder_embed_dim
 __UpperCamelCase				      =  0.02
 __UpperCamelCase				      =  fs_config.encoder_ffn_embed_dim
 __UpperCamelCase				      =  1E-5
 __UpperCamelCase				      =  0.0
 __UpperCamelCase				      =  fs_config.encoder_attention_heads
 __UpperCamelCase				      =  fs_config.conv_pos_groups
 __UpperCamelCase				      =  fs_config.conv_pos
 __UpperCamelCase				      =  len(__A       )
 __UpperCamelCase				      =  fs_config.encoder_layers

 return config







@torch.no_grad()
def _lowercase   (  __A							,__A=None       ):



 '''simple docstring'''



 __UpperCamelCase				      =  distilhubert().model.model

 if config_path is not None:
  __UpperCamelCase				      =  HubertConfig.from_pretrained(__A       )
 else:
  __UpperCamelCase				      =  convert_config(__A       )
 __UpperCamelCase				      =  model.eval()

 __UpperCamelCase				      =  WavaVecaFeatureExtractor(
     feature_size=1							,sampling_rate=16_000							,padding_value=0							,do_normalize=__A							,return_attention_mask=__A							,)
 __UpperCamelCase				      =  HubertModel(__A       )

 recursively_load_weights(__A							,__A       )

 feature_extractor.save_pretrained(__A       )
 hf_model.save_pretrained(__A       )


if __name__ == "__main__":
 a__						:			List[str]										=			argparse.ArgumentParser()
 parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
 parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
 a__						:			Optional[Any]										=			parser.parse_args()
 convert_hubert_checkpoint(args.pytorch_dump_folder_path, args.config_path)





