code: 
  import argparse
  
  # utilities to import the model weights and config file
  import os
  import pickle as pkl
  
  # PyTorch + new model classes
  import torch
  from torch import nn
  
  from transformers import AutoTokenizer, MegaConfig, MegaForMaskedLM
  
  
  # import the EncoderLayer class used to pretrain
  # !! NOTE !! this requires the version of fairseq that is built when you install the Mega source
  try:
  							from fairseq.modules.mega_layer import MegaEncoderLayer
  except ImportError:
  							raise ImportError('You need to install the version of fairseq from the Mega repo!')
  
  
  
  
  
  class 						lowerCamelCase			(					nn.Module				):
  
  
  
  
  
  
  
  	def __init__(    self			:int      ,    lowercase			:Any      ,    lowercase			:Union[str, Any]      ,    lowercase			:Optional[int] )		->			str:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				super().__init__()
  				SCREAMING_SNAKE_CASE    	=   mega_args
  				SCREAMING_SNAKE_CASE    	=   nn.Embedding(lowercase      ,    self.mega_args.encoder_embed_dim )
  				SCREAMING_SNAKE_CASE    	=   nn.ModuleList([MegaEncoderLayer(self.mega_args ) for _ in range(lowercase )] )
  				SCREAMING_SNAKE_CASE    	=   depth
  
  
  
  
  
  
  
  	def    snake_case__    (    self			:str      ,    lowercase			:str      ,    lowercase			:Any      ,    lowercase			:List[str]=True      ,    lowercase			:Union[str, Any]=0 )		->			int:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				if batch_first:
  							SCREAMING_SNAKE_CASE    	=   input_ids.T
  
  				# to make things more confusing, Mega expects the attention mask to
  				# be (batch, time), but with values of 0 (normal token) and 1 (ignore token)
  				# which is the opposite of what HF returns
  				if ignore_mask_value == 0:
  							SCREAMING_SNAKE_CASE    	=   1 - attention_mask
  
  				# get token embeddings from IDs
  				SCREAMING_SNAKE_CASE    	=   self.embedding_layer(lowercase )
  
  				# pass through the Mega layers
  				# input is (time, batch, encoder dim) and output is the same
  				for encoder in self.encoders:
  							SCREAMING_SNAKE_CASE    	=   encoder(lowercase      ,    lowercase )
  
  				# return according to the shape specified
  				if batch_first:
  							# (T, B, H) --> (B, T, H)
  							return torch.transpose(lowercase      ,    0      ,    1 )
  				else:
  							return embeds
  
  
  
  
  
  class 						lowerCamelCase			(					nn.Module				):
  
  
  
  
  
  
  
  	def __init__(    self			:Union[str, Any]      ,    lowercase			:Tuple      ,    lowercase			:Optional[Any]      ,    lowercase			:Dict )		->			Optional[Any]:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				super().__init__()
  				SCREAMING_SNAKE_CASE    	=   MegaLM(lowercase      ,    lowercase      ,    lowercase )
  				SCREAMING_SNAKE_CASE    	=   nn.Linear(mega_args.encoder_embed_dim      ,    lowercase )
  				SCREAMING_SNAKE_CASE    	=   nn.Dropout(p=0.1 )
  
  
  
  
  
  
  
  	def    snake_case__    (    self			:List[str]      ,    lowercase			:Tuple      ,    lowercase			:int      ,    lowercase			:Union[str, Any]=True      ,    lowercase			:List[str]=0 )		->			Any:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				SCREAMING_SNAKE_CASE    	=   self.mega(lowercase      ,    lowercase      ,    lowercase      ,    lowercase )
  				return self.mlm_head(self.dropout(lowercase ) )
  
  
  
  def       a			(   a		,	a		,	a  )    ->Optional[Any]:
  
  
  
  
  			'''simple docstring'''
  
  
  			with open(os.path.join(a		,	'''model_args.pkl'''  )		,	'''rb'''  ) as f:
  						SCREAMING_SNAKE_CASE    	=   pkl.load(a  )
  
  			# load the original encoder
  			SCREAMING_SNAKE_CASE    	=   OriginalMegaForMaskedLM(**a  ).eval()
  
  			# load its weights
  			print(
  			    '''Original Mega encoder:'''		,	original_mlm.mega.load_state_dict(
  			        torch.load(os.path.join(a		,	'''encoder_weights.pt'''  )		,	map_location='''cpu'''  )  )		,	)
  			print(
  			    '''Original Mega MLM layer:'''		,	original_mlm.mlm_head.load_state_dict(
  			        torch.load(os.path.join(a		,	'''mlm_head_weights.pt'''  )		,	map_location='''cpu'''  )  )		,	)
  
  			# create a new config from the old one
  			SCREAMING_SNAKE_CASE    	=   MegaConfig(
  			    num_hidden_layers=mega_original_args['''depth''']		,	vocab_size=mega_original_args['''vocab_size''']		,	hidden_size=mega_original_args['''mega_args'''].encoder_embed_dim		,	shared_representation_size=mega_original_args['''mega_args'''].encoder_z_dim		,	intermediate_size=mega_original_args['''mega_args'''].encoder_hidden_dim		,	ema_projection_size=mega_original_args['''mega_args'''].encoder_n_dim		,	dropout_prob=mega_original_args['''mega_args'''].dropout		,	attention_probs_dropout_prob=mega_original_args['''mega_args'''].attention_dropout		,	hidden_dropout_prob=mega_original_args['''mega_args'''].hidden_dropout		,	activation=mega_original_args['''mega_args'''].activation_fn		,	attention_activation=mega_original_args['''mega_args'''].attention_activation_fn		,	bidirectional=mega_original_args['''mega_args'''].bidirectional		,	use_chunking=mega_original_args['''mega_args'''].encoder_chunk_size > 0		,	chunk_size=mega_original_args['''mega_args'''].encoder_chunk_size		,	truncation=mega_original_args['''mega_args'''].truncation_length		,	normalization_type=mega_original_args['''mega_args'''].normalization_type		,	normalize_before_mega=a		,	norm_affine=a		,	use_feature_dropout=mega_original_args['''mega_args'''].feature_dropout		,	relative_positional_bias=mega_original_args['''mega_args'''].rel_pos_bias		,	max_positions=mega_original_args['''mega_args'''].max_source_positions		,	nffn_hidden_size=mega_original_args['''mega_args'''].encoder_ffn_embed_dim		,	normalize_before_ffn=mega_original_args['''mega_args'''].normalize_before		,	nffn_activation_dropout_prob=0.0		,	add_token_type_embeddings=a		,	add_lm_hidden_dense_layer=a		,	)
  
  			SCREAMING_SNAKE_CASE    	=   MegaForMaskedLM(a  ).eval()
  
  			# the originl checkpoint just uses nn.Embedding for the word embeddings
  			# we use a wrapper module for embeddings to add support for positional embeddings
  			SCREAMING_SNAKE_CASE    	=   original_mlm.mega.embedding_layer.weight
  
  			# modify the state dictionary of the original checkpoint to account for naming issues in the Hugging Face
  			# ecosystem -- any names containing "beta" or "gamma" aren't safe to use and are renamed upon _load_pretrained,
  			# also renaming previously confusing parameter names
  			SCREAMING_SNAKE_CASE    	=   original_mlm.mega.encoders.state_dict()
  			SCREAMING_SNAKE_CASE    	=   {}
  			for module_name in original_state_dict.keys():
  						SCREAMING_SNAKE_CASE    	=   None
  						# have to handle gamma, beta, and alpha differently due to their use
  						# in multiple modules within the original repository;
  						# beta is used in EMA, MovingAverageGatedAttention, and RotaryRelativePositionalBias, and must be renamed due to flax/tf weights
  						# the EMA sublayer was renamed from "move" to "ema_gate" for readability, so that is also done here
  						if "beta" in module_name:
  									# EMA sub-layers were always called "move" in the original repo
  									if "move.beta" in module_name:
  												SCREAMING_SNAKE_CASE    	=   module_name.replace('''move.beta'''		,	'''ema_gate.ema_expansion_matrix'''  )
  									elif "mega_layer.beta" in module_name:
  												SCREAMING_SNAKE_CASE    	=   module_name.replace('''beta'''		,	'''qk_bias'''  )
  									else:
  												SCREAMING_SNAKE_CASE    	=   module_name.replace('''beta'''		,	'''b_param'''  )
          # beta is used in EMA and MovingAverageGatedAttention, and must be renamed due to flax/tf weights
  						elif "gamma" in module_name:
  									if "move.gamma" in module_name:
  												SCREAMING_SNAKE_CASE    	=   module_name.replace('''move.gamma'''		,	'''ema_gate.kernel_projection_matrix'''  )
  									elif "mega_layer.gamma" in module_name:
  												SCREAMING_SNAKE_CASE    	=   module_name.replace('''gamma'''		,	'''qk_weight'''  )
  									else:
  												SCREAMING_SNAKE_CASE    	=   module_name.replace('''gamma'''		,	'''g_param'''  )
          # alpha is used in EMA and positional bias; renaming to improve readability
  						elif "move.alpha" in module_name:
  									SCREAMING_SNAKE_CASE    	=   module_name.replace('''move.alpha'''		,	'''ema_gate.decay_factor'''  )
  						# delta is only used in EMA; renaming to improve readability
  						elif "move.delta" in module_name:
  									SCREAMING_SNAKE_CASE    	=   module_name.replace('''move.delta'''		,	'''ema_gate.damping_factor'''  )
  						# omega is only used in EMA; renaming to improve readability
  						elif "omega" in module_name:
  									SCREAMING_SNAKE_CASE    	=   module_name.replace('''move.omega'''		,	'''ema_gate.residual_weight'''  )
  
  						if new_module_name:
  									SCREAMING_SNAKE_CASE    	=   new_module_name
  
  			if len(a  ) != 0:
  						print(F"""Renaming these keys: {updated_keys.keys()}"""  )
  			else:
  						print('''No need to rename state dict entries'''  )
  			for old, new in updated_keys.items():
  						SCREAMING_SNAKE_CASE    	=   original_state_dict.pop(a  )
  
  			# now attempt to load the state dictionary with updated names
  			# note that we now call it `mega.layers` instead of `mega.encoders` due to hugging face style
  			print('''HF Mega encoder:'''		,	hf_mlm.mega.layers.load_state_dict(a  )  )
  
  			# load the MLM head weights directly
  			print(
  			    '''HF Mega MLM layer:'''		,	hf_mlm.mlm_head.load_state_dict(
  			        torch.load(os.path.join(a		,	'''mlm_head_weights.pt'''  )		,	map_location='''cpu'''  )  )		,	)
  
  			# test on a randomly generated input sequence
  			SCREAMING_SNAKE_CASE    	=   torch.randint(0		,	hf_config.vocab_size		,	size=(4, 256)  )
  			SCREAMING_SNAKE_CASE    	=   torch.ones_like(a  )
  			# mask a few tokens to make sure masking is applied appropriately :)
  			SCREAMING_SNAKE_CASE    	=   0
  
  			# run forward passes
  			SCREAMING_SNAKE_CASE    	=   original_mlm(a		,	a		,	batch_first=a		,	ignore_mask_value=0  )
  			SCREAMING_SNAKE_CASE    	=   hf_mlm(a		,	a  )[0]
  
  			# print shapes and diff
  			print(F"""original output {original_output.shape}"""  )
  			print(F"""hf output {hf_output.shape}"""  )
  			print(F"""max diff: {(original_output - hf_output).max()}"""  )  # 0.0
  			SCREAMING_SNAKE_CASE    	=   torch.allclose(a		,	a		,	atol=1E-3  )
  
  			if success:
  						print('''Yay!'''  )
  						hf_mlm.save_pretrained(a  )
  			else:
  						raise RuntimeError(F"""Something's broken :(\nOriginal:\n{original_output}\n\nHF\n{hf_output}\n{hf_mlm}"""  )
  
  			if includes_tokenizer:
  						print('''Transferring tokenizer'''  )
  						SCREAMING_SNAKE_CASE    	=   AutoTokenizer.from_pretrained(a  )
  						tokenizer.save_pretrained(a  )
  
  
  if __name__ == "__main__":
  							__lowerCAmelCase				      =					argparse.ArgumentParser()
  
  							parser.add_argument(
  							    '--pretrained_checkpoint_path',
  							    default=None,
  							    type=str,
  							    required=True,
  							    help='Point to the directory containing your model weights using the official Mega repo',
  							)
  
  							parser.add_argument(
  							    '--output_path', default=None, type=str, required=True, help='Location to save the Hugging Face version'
  							)
  
  							parser.add_argument(
  							    '--includes_tokenizer',
  							    action='store_true',
  							    help='Use this flag if there is a Hugging Face tokenizer in the original checkpoint repo',
  							)
  
  							__lowerCAmelCase				      =					parser.parse_args()
  
  							convert_checkpoint_to_huggingface(args.pretrained_checkpoint_path, args.output_path, args.includes_tokenizer)
style_context: 
  import argparse
  import os
  from dataclasses import dataclass
  
  import torch
  import torchaudio
  from tqdm import tqdm
  
  from transformers import AutoConfig, AutoModelForCTC, AutoProcessor
  
  
  
  
  
  class 						lowerCamelCase			:
  
  
  
  
  
  
  
  	def __init__(    self			:Optional[int]      ,    lowercase			:Optional[Any]      ,    lowercase			:Optional[int]      ,    lowercase			:int )		->			Optional[int]:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				SCREAMING_SNAKE_CASE    	=   cuda
  				SCREAMING_SNAKE_CASE    	=   AutoConfig.from_pretrained(lowercase )
  				SCREAMING_SNAKE_CASE    	=   AutoModelForCTC.from_pretrained(lowercase )
  				self.model.eval()
  				if self.cuda:
  							self.model.to(device='''cuda''' )
  				SCREAMING_SNAKE_CASE    	=   AutoProcessor.from_pretrained(lowercase )
  				SCREAMING_SNAKE_CASE    	=   torchaudio.transforms.Resample(lowercase      ,    1_6_0_0_0 )
  				SCREAMING_SNAKE_CASE    	=   0
  				SCREAMING_SNAKE_CASE    	=   list(self.processor.tokenizer.get_vocab().keys() )
  				for i in range(len(lowercase ) ):
  							if vocab[i] == "[PAD]" or vocab[i] == "<pad>":
  										SCREAMING_SNAKE_CASE    	=   i
  				print('''Blank Token id [PAD]/<pad>'''      ,    lowercase )
  				SCREAMING_SNAKE_CASE    	=   blank_id
  
  
  
  
  
  
  
  	def    snake_case__    (    self			:Any      ,    lowercase			:int )		->			Optional[int]:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				SCREAMING_SNAKE_CASE	,       SCREAMING_SNAKE_CASE    	=   torchaudio.load(lowercase )
  				SCREAMING_SNAKE_CASE    	=   self.resampler(lowercase ).squeeze().numpy()
  				return speech
  
  
  
  
  
  
  
  	def    snake_case__    (    self			:Union[str, Any]      ,    lowercase			:Optional[Any] )		->			int:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				SCREAMING_SNAKE_CASE    	=   self.blank_id
  				SCREAMING_SNAKE_CASE    	=   '''|'''.join(item['''sent'''].split(''' ''' ) )
  				if not os.path.isfile(item['''wav_path'''] ):
  							print(item['''wav_path''']      ,    '''not found in wavs directory''' )
  
  				SCREAMING_SNAKE_CASE    	=   self.speech_file_to_array_fn(item['''wav_path'''] )
  				SCREAMING_SNAKE_CASE    	=   self.processor(lowercase      ,    sampling_rate=1_6_0_0_0      ,    return_tensors='''pt'''      ,    padding=lowercase )
  				if self.cuda:
  							SCREAMING_SNAKE_CASE    	=   inputs.to(device='''cuda''' )
  
  				with torch.no_grad():
  							SCREAMING_SNAKE_CASE    	=   self.model(inputs.input_values ).logits
  
  				# get the emission probability at frame level
  				SCREAMING_SNAKE_CASE    	=   torch.log_softmax(lowercase      ,    dim=-1 )
  				SCREAMING_SNAKE_CASE    	=   emissions[0].cpu().detach()
  
  				# get labels from vocab
  				SCREAMING_SNAKE_CASE    	=   ([''''''] + list(self.processor.tokenizer.get_vocab().keys() ))[
  				    :-1
  				]  # logits don't align with the tokenizer's vocab
  
  				SCREAMING_SNAKE_CASE    	=   {c: i for i, c in enumerate(lowercase )}
  				SCREAMING_SNAKE_CASE    	=   []
  				for c in transcript:
  							if c in dictionary:
  										tokens.append(dictionary[c] )
  
  				def get_trellis(lowercase			:List[Any]      ,    lowercase			:Any      ,    lowercase			:int=0 ):
  							SCREAMING_SNAKE_CASE    	=   emission.size(0 )
  							SCREAMING_SNAKE_CASE    	=   len(lowercase )
  
  							# Trellis has extra diemsions for both time axis and tokens.
  							# The extra dim for tokens represents <SoS> (start-of-sentence)
  							# The extra dim for time axis is for simplification of the code.
  							SCREAMING_SNAKE_CASE    	=   torch.full((num_frames + 1, num_tokens + 1)      ,    -float('''inf''' ) )
  							SCREAMING_SNAKE_CASE    	=   0
  							for t in range(lowercase ):
  										SCREAMING_SNAKE_CASE    	=   torch.maximum(
  										    # Score for staying at the same token
  										    trellis[t, 1:] + emission[t, blank_id]      ,    trellis[t, :-1] + emission[t, tokens]      ,    )
  							return trellis
  
  				SCREAMING_SNAKE_CASE    	=   get_trellis(lowercase      ,    lowercase      ,    lowercase )
  
  
  
  
  
  				@dataclass
  				class 						lowerCamelCase			:
  					UpperCamelCase_     :				int
  					UpperCamelCase_     :				int
  					UpperCamelCase_     :				float
  
  				def backtrack(lowercase			:Optional[int]      ,    lowercase			:List[str]      ,    lowercase			:Union[str, Any]      ,    lowercase			:Optional[Any]=0 ):
  							SCREAMING_SNAKE_CASE    	=   trellis.size(1 ) - 1
  							SCREAMING_SNAKE_CASE    	=   torch.argmax(trellis[:, j] ).item()
  
  							SCREAMING_SNAKE_CASE    	=   []
  							for t in range(lowercase      ,    0      ,    -1 ):
  										# 1. Figure out if the current position was stay or change
  										# Note (again):
  										# `emission[J-1]` is the emission at time frame `J` of trellis dimension.
  										# Score for token staying the same from time frame J-1 to T.
  										SCREAMING_SNAKE_CASE    	=   trellis[t - 1, j] + emission[t - 1, blank_id]
  										# Score for token changing from C-1 at T-1 to J at T.
  										SCREAMING_SNAKE_CASE    	=   trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]
  
  										# 2. Store the path with frame-wise probability.
  										SCREAMING_SNAKE_CASE    	=   emission[t - 1, tokens[j - 1] if changed > stayed else 0].exp().item()
  										# Return token index and time index in non-trellis coordinate.
  										path.append(Point(j - 1      ,    t - 1      ,    lowercase ) )
  
  										# 3. Update the token
  										if changed > stayed:
  													j -= 1
  													if j == 0:
  																break
  							else:
  										raise ValueError('''Failed to align''' )
  							return path[::-1]
  
  				SCREAMING_SNAKE_CASE    	=   backtrack(lowercase      ,    lowercase      ,    lowercase      ,    lowercase )
  
  
  
  
  
  				@dataclass
  				class 						lowerCamelCase			:
  					UpperCamelCase_     :				str
  					UpperCamelCase_     :				int
  					UpperCamelCase_     :				int
  					UpperCamelCase_     :				float
  
  
  
  
  
  
  
  					def __repr__(    self			:Optional[Any] )		->			str:
  
  
  
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return f"""{self.label}\t{self.score:4.2f}\t{self.start*2_0:5d}\t{self.end*2_0:5d}"""
  
  
  
  
  
  
  
  					@property
  					def    snake_case__    (    self			:List[str] )		->			Any:
  
  
  
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.end - self.start
  
  				def merge_repeats(lowercase			:List[str] ):
  							SCREAMING_SNAKE_CASE	,       SCREAMING_SNAKE_CASE    	=   0, 0
  							SCREAMING_SNAKE_CASE    	=   []
  							while ia < len(lowercase ):
  										while ia < len(lowercase ) and path[ia].token_index == path[ia].token_index:
  													ia += 1
  										SCREAMING_SNAKE_CASE    	=   sum(path[k].score for k in range(lowercase      ,    lowercase ) ) / (ia - ia)
  										segments.append(
  										    Segment(
  										        transcript[path[ia].token_index]      ,    path[ia].time_index      ,    path[ia - 1].time_index + 1      ,    lowercase      ,    ) )
  										SCREAMING_SNAKE_CASE    	=   ia
  							return segments
  
  				SCREAMING_SNAKE_CASE    	=   merge_repeats(lowercase )
  				with open(item['''out_path''']      ,    '''w''' ) as out_align:
  							for seg in segments:
  										out_align.write(str(lowercase ) + '''\n''' )
  
  
  
  
  
  
  
  	def    snake_case__    (    self			:List[Any]      ,    lowercase			:Dict      ,    lowercase			:Tuple      ,    lowercase			:str )		->			Dict:
  
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  				if not os.path.exists(lowercase ):
  							os.makedirs(lowercase )
  
  				# load text file
  				SCREAMING_SNAKE_CASE    	=   open(lowercase      ,    encoding='''utf8''' ).readlines()
  
  				SCREAMING_SNAKE_CASE    	=   []
  				for line in lines:
  							if len(line.strip().split('''\t''' ) ) != 2:
  										print('''Script must be in format: 00001  this is my sentence''' )
  										exit()
  
  							SCREAMING_SNAKE_CASE	,       SCREAMING_SNAKE_CASE    	=   line.strip().split('''\t''' )
  							SCREAMING_SNAKE_CASE    	=   os.path.join(lowercase      ,    wav_name + '''.wav''' )
  							SCREAMING_SNAKE_CASE    	=   os.path.join(lowercase      ,    wav_name + '''.txt''' )
  
  							items.append({'''sent''': sentence, '''wav_path''': wav_path, '''out_path''': out_path} )
  				print('''Number of samples found in script file'''      ,    len(lowercase ) )
  
  				for item in tqdm(lowercase ):
  							self.align_single_sample(lowercase )
  
  
  
  def       a			(   )    ->Optional[int]:
  
  
  
  
  			'''simple docstring'''
  
  
  			SCREAMING_SNAKE_CASE    	=   argparse.ArgumentParser()
  
  			parser.add_argument(
  			    '''--model_name'''		,	type=a		,	default='''arijitx/wav2vec2-xls-r-300m-bengali'''		,	help='''wav2vec model name'''  )
  			parser.add_argument('''--wav_dir'''		,	type=a		,	default='''./wavs'''		,	help='''directory containing wavs'''  )
  			parser.add_argument('''--text_file'''		,	type=a		,	default='''script.txt'''		,	help='''file containing text'''  )
  			parser.add_argument('''--input_wavs_sr'''		,	type=a		,	default=1_6000		,	help='''sampling rate of input audios'''  )
  			parser.add_argument(
  			    '''--output_dir'''		,	type=a		,	default='''./out_alignment'''		,	help='''output directory containing the alignment files'''  )
  			parser.add_argument('''--cuda'''		,	action='''store_true'''  )
  
  			SCREAMING_SNAKE_CASE    	=   parser.parse_args()
  
  			SCREAMING_SNAKE_CASE    	=   WavaVecaAligner(args.model_name		,	args.input_wavs_sr		,	args.cuda  )
  			aligner.align_data(args.wav_dir		,	args.text_file		,	args.output_dir  )
  
  
  if __name__ == "__main__":
  							main()