



"""simple docstring"""






import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
   from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse('''1.6'''):
   lowerCAmelCase__        =				True
   from torch.cuda.amp import autocast


lowerCAmelCase__        =				logging.getLogger(__name__)






@dataclass
class        __snake_case       :
  snake_case__				:   str  			=						field(
      metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"})
  snake_case__				:   Optional[str]  			=						field(
      default=_lowercase  ,     metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}  ,     )
  snake_case__				:   Optional[bool]  			=						field(
      default=_lowercase  ,     metadata={"help": "Whether to freeze the feature extractor layers of the model."})
  snake_case__				:   Optional[bool]  			=						field(
      default=_lowercase  ,     metadata={"help": "Whether to log verbose messages or not."}  ,     )
  snake_case__				:   Optional[float]  			=						field(
      default=2.0  ,     metadata={"help": "Maximum temperature for gumbel softmax."})
  snake_case__				:   Optional[float]  			=						field(
      default=0.5  ,     metadata={"help": "Minimum temperature for gumbel softmax."})
  snake_case__				:   Optional[float]  			=						field(
      default=0.999_995  ,     metadata={"help": "Decay of gumbel temperature during training."})






def        snake_case_  (      A_					:		ModelArguments,				A_					:		TrainingArguments						):


       '''simple docstring'''





       logging.basicConfig(
           format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''',				datefmt='''%m/%d/%Y %H:%M:%S''',				handlers=[logging.StreamHandler(sys.stdout						)],				)
       _lowerCamelCase						:				List[str]    			= logging.WARNING
       if model_args.verbose_logging:
              _lowerCamelCase						:				Optional[Any]    			= logging.DEBUG
       elif trainer_utils.is_main_process(training_args.local_rank						):
              _lowerCamelCase						:				int    			= logging.INFO
       logger.setLevel(A_						)






@dataclass
class        __snake_case       :
  snake_case__				:   str  			=						field(
      default=_lowercase  ,     metadata={"help": "The name of the dataset to use (via the datasets library)."})
  snake_case__				:   Optional[str]  			=						field(
      default=_lowercase  ,     metadata={"help": "The configuration name of the dataset to use (via the datasets library)."})
  snake_case__				:   Optional[str]  			=						field(
      default="train"  ,     metadata={
          "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
      }  ,     )
  snake_case__				:   Optional[str]  			=						field(
      default="validation"  ,     metadata={
          "help": (
              "The name of the validation data set split to use (via the datasets library). Defaults to 'validation'"
          )
      }  ,     )
  snake_case__				:   Optional[str]  			=						field(
      default="file"  ,     metadata={"help": "Column in the dataset that contains speech file path. Defaults to 'file'"}  ,     )
  snake_case__				:   bool  			=						field(
      default=_lowercase  ,     metadata={"help": "Overwrite the cached preprocessed datasets or not."})
  snake_case__				:   Optional[int]  			=						field(
      default=1  ,     metadata={
          "help": "The percentage of the train set used as validation set in case there's no validation split"
      }  ,     )
  snake_case__				:   Optional[int]  			=						field(
      default=_lowercase  ,     metadata={"help": "The number of processes to use for the preprocessing."}  ,     )
  snake_case__				:   Optional[float]  			=						field(
      default=20.0  ,     metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"})






@dataclass
class        __snake_case       :
  snake_case__				:   WavaVecaForPreTraining
  snake_case__				:   WavaVecaFeatureExtractor
  snake_case__				:   Union[bool, str]  			=						"longest"
  snake_case__				:   Optional[int]  			=						None
  snake_case__				:   Optional[int]  			=						None




  def __call__( self     :						int						,					__lowerCAmelCase     :						List[Dict[str, Union[List[int], torch.Tensor]]]		):






         """simple docstring"""






         _lowerCamelCase						:				Tuple    			= self.feature_extractor.pad(
             __lowerCAmelCase						,					max_length=self.max_length						,					padding=self.padding						,					pad_to_multiple_of=self.pad_to_multiple_of						,					return_tensors='''pt'''						,					)
         _lowerCamelCase						:				int    			= self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]		)

         _lowerCamelCase						:				str    			= batch['''input_values'''].shape[0]

         # make sure that no loss is computed on padded inputs
         if batch["attention_mask"] is not None:
                # compute real output lengths according to convolution formula
                _lowerCamelCase						:				Any    			= self.model._get_feat_extract_output_lengths(batch['''attention_mask'''].sum(-1		)		).to(
                    torch.long		)

                _lowerCamelCase						:				Any    			= torch.zeros(
                    (batch_size, mask_indices_seq_length)						,					dtype=torch.long						,					device=batch['''input_values'''].device		)

                # these two operations makes sure that all values
                # before the output lengths indices are attended to
                _lowerCamelCase						:				Union[str, Any]    			= 1
                _lowerCamelCase						:				List[Any]    			= attention_mask.flip([-1]		).cumsum(-1		).flip([-1]		).bool()

         # sample randomly masked indices
         _lowerCamelCase						:				Optional[int]    			= _compute_mask_indices(
             (batch_size, mask_indices_seq_length)						,					self.model.config.mask_time_prob						,					self.model.config.mask_time_length						,					attention_mask=__lowerCAmelCase						,					min_masks=2						,					)

         return batch






class        __snake_case       (	_lowercase):




  def __init__( self     :						int						,					*__lowerCAmelCase     :						Optional[Any]						,					__lowerCAmelCase     :						List[Any]=1						,					__lowerCAmelCase     :						str=0						,					__lowerCAmelCase     :						List[Any]=1.0						,					**__lowerCAmelCase     :						List[Any]		):






         """simple docstring"""






         super().__init__(*__lowerCAmelCase						,					**__lowerCAmelCase		)
         _lowerCamelCase						:				Dict    			= 0
         _lowerCamelCase						:				List[str]    			= max_gumbel_temp
         _lowerCamelCase						:				Tuple    			= min_gumbel_temp
         _lowerCamelCase						:				Union[str, Any]    			= gumbel_temp_decay




  def 						SCREAMING_SNAKE_CASE				( self     :						Any						,					__lowerCAmelCase     :						nn.Module						,					__lowerCAmelCase     :						Dict[str, Union[torch.Tensor, Any]]		):






         """simple docstring"""






         model.train()
         _lowerCamelCase						:				Dict    			= self._prepare_inputs(__lowerCAmelCase		)

         if self.use_amp:
                with autocast():
                       _lowerCamelCase						:				Optional[Any]    			= self.compute_loss(__lowerCAmelCase						,					__lowerCAmelCase		)
         else:
                _lowerCamelCase						:				Dict    			= self.compute_loss(__lowerCAmelCase						,					__lowerCAmelCase		)

         if self.args.n_gpu > 1 or self.deepspeed:
                if model.module.config.ctc_loss_reduction == "mean":
                       _lowerCamelCase						:				List[str]    			= loss.mean()
                elif model.module.config.ctc_loss_reduction == "sum":
                       _lowerCamelCase						:				List[str]    			= loss.sum() / (inputs['''mask_time_indices''']).sum()
                else:
                       raise ValueError(f'''{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']'''		)

         if self.args.gradient_accumulation_steps > 1:
                _lowerCamelCase						:				Optional[int]    			= loss / self.args.gradient_accumulation_steps

         if self.use_amp:
                self.scaler.scale(__lowerCAmelCase		).backward()
         elif self.use_apex:
                with amp.scale_loss(__lowerCAmelCase						,					self.optimizer		) as scaled_loss:
                       scaled_loss.backward()
         elif self.deepspeed:
                self.deepspeed.backward(__lowerCAmelCase		)
         else:
                loss.backward()

         self.num_update_step += 1
         # make sure gumbel softmax temperature is decayed
         if self.args.n_gpu > 1 or self.deepspeed:
                model.module.set_gumbel_temperature(
                    max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step						,					self.min_gumbel_temp		)		)
         else:
                model.set_gumbel_temperature(
                    max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step						,					self.min_gumbel_temp		)		)

         return loss.detach()






def        snake_case_  (      ):


       '''simple docstring'''





       _lowerCamelCase						:				Any    			= HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)

       _lowerCamelCase			,    _lowerCamelCase			,    _lowerCamelCase						:				int    			= parser.parse_args_into_dataclasses()
       configure_logger(A_,				A_						)

       # Downloading and loading a dataset from the hub.
       _lowerCamelCase						:				List[Any]    			= load_dataset(data_args.dataset_name,				data_args.dataset_config_name,				cache_dir=model_args.cache_dir						)

       if "validation" not in datasets.keys():
              # make sure only "validation" and "train" keys remain"
              _lowerCamelCase						:				Dict    			= DatasetDict()
              _lowerCamelCase						:				Union[str, Any]    			= load_dataset(
                  data_args.dataset_name,				data_args.dataset_config_name,				split=F'''{data_args.train_split_name}[:{data_args.validation_split_percentage}%]''',				cache_dir=model_args.cache_dir,				)
              _lowerCamelCase						:				str    			= load_dataset(
                  data_args.dataset_name,				data_args.dataset_config_name,				split=F'''{data_args.train_split_name}[{data_args.validation_split_percentage}%:]''',				cache_dir=model_args.cache_dir,				)
       else:
              # make sure only "validation" and "train" keys remain"
              _lowerCamelCase						:				int    			= DatasetDict()
              _lowerCamelCase						:				Any    			= load_dataset(
                  data_args.dataset_name,				data_args.dataset_config_name,				split='''validation''',				cache_dir=model_args.cache_dir,				)
              _lowerCamelCase						:				Union[str, Any]    			= load_dataset(
                  data_args.dataset_name,				data_args.dataset_config_name,				split=F'''{data_args.train_split_name}''',				cache_dir=model_args.cache_dir,				)

       # only normalized-inputs-training is supported
       _lowerCamelCase						:				int    			= WavaVecaFeatureExtractor.from_pretrained(
           model_args.model_name_or_path,				cache_dir=model_args.cache_dir,				do_normalize=A_						)

       def prepare_dataset(A_					:		Optional[int]						):
              # check that all files have the correct sampling rate
              _lowerCamelCase			,    _lowerCamelCase						:				Dict    			= librosa.load(batch[data_args.speech_file_column],				sr=feature_extractor.sampling_rate						)
              return batch

       # load audio files into numpy arrays
       _lowerCamelCase						:				Tuple    			= datasets.map(
           A_,				num_proc=data_args.preprocessing_num_workers,				remove_columns=datasets['''train'''].column_names						)

       # filter audio files that are too long
       _lowerCamelCase						:				Optional[int]    			= vectorized_datasets.filter(
           lambda A_						: len(data['''speech''']						) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate						)						)

       def normalize(A_					:		Optional[int]						):
              return feature_extractor(batch['''speech'''],				sampling_rate=feature_extractor.sampling_rate						)

       # normalize and transform to `BatchFeatures`
       _lowerCamelCase						:				str    			= vectorized_datasets.map(
           A_,				batched=A_,				num_proc=data_args.preprocessing_num_workers,				load_from_cache_file=not data_args.overwrite_cache,				remove_columns=vectorized_datasets['''train'''].column_names,				)

       # pretraining is only supported for "newer" stable layer norm architecture
       # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
       _lowerCamelCase						:				List[str]    			= WavaVecaConfig.from_pretrained(
           model_args.model_name_or_path,				cache_dir=model_args.cache_dir,				gradient_checkpointing=training_args.gradient_checkpointing,				)

       if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
              raise ValueError(
                  '''PreTraining is only supported for ``config.do_stable_layer_norm=True`` and'''
                  ''' ``config.feat_extract_norm=\'layer\''''						)

       _lowerCamelCase						:				List[Any]    			= WavaVecaForPreTraining(A_						)

       _lowerCamelCase						:				Union[str, Any]    			= DataCollatorForWavaVecaPretraining(model=A_,				feature_extractor=A_						)

       _lowerCamelCase						:				Optional[Any]    			= WavaVecaPreTrainer(
           model=A_,				data_collator=A_,				args=A_,				train_dataset=vectorized_datasets['''train'''],				eval_dataset=vectorized_datasets['''validation'''],				tokenizer=A_,				max_gumbel_temp=model_args.max_gumbel_temperature,				min_gumbel_temp=model_args.min_gumbel_temperature,				gumbel_temp_decay=model_args.gumbel_temperature_decay,				)
       trainer.train()


if __name__ == "__main__":
   main()




import argparse
import json
import os
import re
import shutil

import torch

from transformers import BioGptConfig, BioGptForCausalLM
from transformers.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES
from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
from transformers.utils import WEIGHTS_NAME, logging


logging.set_verbosity_warning()

lowercase__		:str   =    2







class 				lowercase		:




							def __init__(				self   ,*,  # begin keyword-only arguments
							    A__="<s>"   ,A__="<pad>"   ,A__="</s>"   ,A__="<unk>"   ,A__=None   ,):
												lowercase	,						lowercase	,						lowercase	,						lowercase							=						bos, unk, pad, eos
												lowercase							=						[]
												lowercase							=						[]
												lowercase							=						{}
												lowercase							=						self.add_symbol(A__)
												lowercase							=						self.add_symbol(A__)
												lowercase							=						self.add_symbol(A__)
												lowercase							=						self.add_symbol(A__)
												if extra_special_symbols:
																	for s in extra_special_symbols:
																						self.add_symbol(A__)
												lowercase							=						len(self.symbols)




							def __eq__(				self   ,A__):
												return self.indices == other.indices




							def __getitem__(				self   ,A__):
												if idx < len(self.symbols):
																	return self.symbols[idx]
												return self.unk_word




							def __len__(				self):
												return len(self.symbols)




							def __contains__(				self   ,A__):
												return sym in self.indices




							@classmethod
							def  A__      (				cls   ,A__):
												lowercase							=						cls()
												d.add_from_file(A__)
												return d




							def  A__      (				self   ,A__   ,A__=1   ,A__=False):
												if word in self.indices and not overwrite:
																	lowercase							=						self.indices[word]
																	lowercase							=						self.count[idx] + n
																	return idx
												else:
																	lowercase							=						len(self.symbols)
																	lowercase							=						idx
																	self.symbols.append(A__)
																	self.count.append(A__)
																	return idx




							def  A__      (				self   ,A__):
												return 0





							def  A__      (				self   ,A__):
												if isinstance(A__   ,A__):
																	try:
																						with open(A__   ,'''r'''   ,encoding='''utf-8''') as fd:
																											self.add_from_file(A__)
																	except FileNotFoundError as fnfe:
																						raise fnfe
																	except UnicodeError:
																						raise Exception('''Incorrect encoding detected in {}, please rebuild the dataset'''.format(A__))
																	return

												lowercase							=						f.readlines()
												lowercase							=						self._load_meta(A__)

												for line in lines[indices_start_line:]:
																	try:
																						lowercase	,						lowercase							=						line.rstrip().rsplit(''' '''   ,1)
																						if field == "#fairseq:overwrite":
																											lowercase							=						True
																											lowercase	,						lowercase							=						line.rsplit(''' '''   ,1)
																						else:
																											lowercase							=						False
																						lowercase							=						int(A__)
																						lowercase							=						line
																						if word in self and not overwrite:
																											raise RuntimeError(
																											    '''Duplicate word found when loading Dictionary: \'{}\'. '''
																											    '''Duplicate words can overwrite earlier ones by adding the '''
																											    '''#fairseq:overwrite flag at the end of the corresponding row '''
																											    '''in the dictionary file. If using the Camembert model, please '''
																											    '''download an updated copy of the model file.'''.format(A__))
																						self.add_symbol(A__   ,n=A__   ,overwrite=A__)
																	except ValueError:
																						raise ValueError('''Incorrect dictionary format, expected \'<token> <cnt> [flags]\'''')



def        UpperCamelCase    (							lowerCAmelCase__		):
					'''simple docstring'''
					# (1) remove word breaking symbol, (2) add word ending symbol where the word is not broken up,
					# e.g.: d = {'le@@': 5, 'tt@@': 6, 'er': 7} => {'le': 5, 'tt': 6, 'er</w>': 7}
					lowercase							=						dict((re.sub(R'''@@$''' ,			'''''' ,			lowerCAmelCase__		), v) if k.endswith('''@@'''		) else (re.sub(R'''$''' ,			'''</w>''' ,			lowerCAmelCase__		), v) for k, v in d.items()		)
					lowercase							=						'''<s> <pad> </s> <unk>'''.split()
					# restore the special tokens
					for k in keep_keys:
										del da[f'{k}</w>']
										lowercase							=						d[k]  # restore
					return da



def        UpperCamelCase    (							lowerCAmelCase__ ,			lowerCAmelCase__		):
					'''simple docstring'''
					# prep
					if not os.path.exists(lowerCAmelCase__		):
										raise ValueError(f'path {biogpt_checkpoint_path} does not exist!'		)
					os.makedirs(lowerCAmelCase__ ,			exist_ok=lowerCAmelCase__		)
					print(f'Writing results to {pytorch_dump_folder_path}'		)

					# handle various types of models

					lowercase							=						os.path.join(lowerCAmelCase__ ,			'''checkpoint.pt'''		)
					if not os.path.isfile(lowerCAmelCase__		):
										raise ValueError(f'path to the file {checkpoint_file} does not exist!'		)
					lowercase							=						torch.load(lowerCAmelCase__ ,			map_location='''cpu'''		)

					lowercase							=						chkpt['''cfg''']['''model''']

					# dicts
					lowercase							=						os.path.join(lowerCAmelCase__ ,			'''dict.txt'''		)
					if not os.path.isfile(lowerCAmelCase__		):
										raise ValueError(f'path to the file {dict_file} does not exist!'		)
					lowercase							=						Dictionary.load(lowerCAmelCase__		)
					lowercase							=						rewrite_dict_keys(src_dict.indices		)
					lowercase							=						len(lowerCAmelCase__		)
					lowercase							=						os.path.join(lowerCAmelCase__ ,			VOCAB_FILES_NAMES['''vocab_file''']		)
					print(f'Generating {src_vocab_file} of {src_vocab_size} records'		)
					with open(lowerCAmelCase__ ,			'''w''' ,			encoding='''utf-8'''		) as f:
										f.write(json.dumps(lowerCAmelCase__ ,			ensure_ascii=lowerCAmelCase__ ,			indent=lowerCAmelCase__		)		)

					# merges_file (bpecodes)
					lowercase							=						os.path.join(lowerCAmelCase__ ,			'''bpecodes'''		)
					if not os.path.isfile(lowerCAmelCase__		):
										raise ValueError(f'path to the file {bpecodes_file} does not exist!'		)

					lowercase							=						os.path.join(lowerCAmelCase__ ,			VOCAB_FILES_NAMES['''merges_file''']		)
					shutil.copyfile(lowerCAmelCase__ ,			lowerCAmelCase__		)

					# model config
					lowercase							=						os.path.join(lowerCAmelCase__ ,			'''config.json'''		)

					lowercase							=						{
					    '''activation_dropout''': args['''activation_dropout'''],
					    '''architectures''': ['''BioGptForCausalLM'''],
					    '''attention_probs_dropout_prob''': args['''attention_dropout'''],
					    '''bos_token_id''': 0,
					    '''eos_token_id''': 2,
					    '''hidden_act''': args['''activation_fn'''],
					    '''hidden_dropout_prob''': args['''dropout'''],
					    '''hidden_size''': args['''decoder_embed_dim'''],
					    '''initializer_range''': 0.02,
					    '''intermediate_size''': args['''decoder_ffn_embed_dim'''],
					    '''layer_norm_eps''': 1E-12,
					    '''layerdrop''': args['''decoder_layerdrop'''],
					    '''max_position_embeddings''': args['''max_target_positions'''],
					    '''model_type''': '''biogpt''',
					    '''num_attention_heads''': args['''decoder_attention_heads'''],
					    '''num_hidden_layers''': args['''decoder_layers'''],
					    '''pad_token_id''': 1,
					    '''scale_embedding''': not args['''no_scale_embedding'''],
					    '''tie_word_embeddings''': args['''share_decoder_input_output_embed'''],
					    '''vocab_size''': src_vocab_size,
					}

					# good hparam defaults to start with

					print(f'Generating {biogpt_model_config_file}'		)
					with open(lowerCAmelCase__ ,			'''w''' ,			encoding='''utf-8'''		) as f:
										f.write(json.dumps(lowerCAmelCase__ ,			ensure_ascii=lowerCAmelCase__ ,			indent=lowerCAmelCase__		)		)

					# tokenizer config
					lowercase							=						os.path.join(lowerCAmelCase__ ,			lowerCAmelCase__		)

					lowercase							=						{
					    '''bos_token''': '''<s>''',
					    '''eos_token''': '''</s>''',
					    '''model_max_length''': 1024,
					    '''pad_token''': '''<pad>''',
					    '''special_tokens_map_file''': None,
					    '''tokenizer_class''': '''BioGptTokenizer''',
					    '''unk_token''': '''<unk>''',
					}

					print(f'Generating {biogpt_tokenizer_config_file}'		)
					with open(lowerCAmelCase__ ,			'''w''' ,			encoding='''utf-8'''		) as f:
										f.write(json.dumps(lowerCAmelCase__ ,			ensure_ascii=lowerCAmelCase__ ,			indent=lowerCAmelCase__		)		)

					# model
					lowercase							=						chkpt['''model''']

					# remove unneeded keys
					lowercase							=						[
					    '''decoder.version''',
					]
					for k in ignore_keys:
										model_state_dict.pop(lowerCAmelCase__ ,			lowerCAmelCase__		)

					lowercase							=						list(model_state_dict.keys()		)
					for layer_name in layer_names:
										if layer_name.endswith('''output_projection.weight'''		):
															lowercase							=						model_state_dict.pop(lowerCAmelCase__		)
										else:
															lowercase							=						model_state_dict.pop(lowerCAmelCase__		)

					lowercase							=						BioGptConfig.from_pretrained(lowerCAmelCase__		)
					lowercase							=						BioGptForCausalLM(lowerCAmelCase__		)

					# check that it loads ok
					model_new.load_state_dict(lowerCAmelCase__		)

					# save
					lowercase							=						os.path.join(lowerCAmelCase__ ,			lowerCAmelCase__		)
					print(f'Generating {pytorch_weights_dump_path}'		)
					torch.save(lowerCAmelCase__ ,			lowerCAmelCase__		)

					print('''Conversion is done!'''		)


if __name__ == "__main__":
							lowercase__		:Union[str, Any]   =    argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    "--biogpt_checkpoint_path",
							    default=None,
							    type=str,
							    required=True,
							    help=(
							        "Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,"
							        " bpecodes, etc."
							    ),
							)
							parser.add_argument(
							    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
							)
							lowercase__		:Any   =    parser.parse_args()
							convert_biogpt_checkpoint_to_pytorch(args.biogpt_checkpoint_path, args.pytorch_dump_folder_path)




