





'''simple docstring'''






import argparse
import os
from dataclasses import dataclass

import torch
import torchaudio
from tqdm import tqdm

from transformers import AutoConfig, AutoModelForCTC, AutoProcessor



class  lowerCamelCase__				:


	"""simple docstring"""





	def __init__(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)       ->				Optional[Any]:
					_A :				List[Any]				=  cuda
					_A :				List[str]				=  AutoConfig.from_pretrained(UpperCAmelCase__			)
					_A :				List[Any]				=  AutoModelForCTC.from_pretrained(UpperCAmelCase__			)
					self.model.eval()
					if self.cuda:
									self.model.to(device='''cuda'''			)
					_A :				Optional[Any]				=  AutoProcessor.from_pretrained(UpperCAmelCase__			)
					_A :				List[str]				=  torchaudio.transforms.Resample(UpperCAmelCase__   ,     1_6_0_0_0			)
					_A :				Union[str, Any]				=  0
					_A :				Optional[int]				=  list(self.processor.tokenizer.get_vocab().keys()			)
					for i in range(len(UpperCAmelCase__			)			):
									if vocab[i] == "[PAD]" or vocab[i] == "<pad>":
													_A :				List[Any]				=  i
					print('''Blank Token id [PAD]/<pad>'''   ,     UpperCAmelCase__			)
					_A :				List[Any]				=  blank_id




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__			)       ->				int:
					_A	,       _A :				Tuple				=  torchaudio.load(UpperCAmelCase__			)
					_A :				Any				=  self.resampler(UpperCAmelCase__			).squeeze().numpy()
					return speech




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__			)       ->				Dict:
					_A :				List[Any]				=  self.blank_id
					_A :				List[str]				=  '''|'''.join(item['''sent'''].split(''' '''			)			)
					if not os.path.isfile(item['''wav_path''']			):
									print(item['''wav_path''']   ,     '''not found in wavs directory'''			)

					_A :				Optional[Any]				=  self.speech_file_to_array_fn(item['''wav_path''']			)
					_A :				Optional[int]				=  self.processor(UpperCAmelCase__   ,     sampling_rate=1_6_0_0_0   ,     return_tensors='''pt'''   ,     padding=UpperCAmelCase__			)
					if self.cuda:
									_A :				int				=  inputs.to(device='''cuda'''			)

					with torch.no_grad():
									_A :				int				=  self.model(inputs.input_values			).logits

					# get the emission probability at frame level
					_A :				Tuple				=  torch.log_softmax(UpperCAmelCase__   ,     dim=-1			)
					_A :				Optional[int]				=  emissions[0].cpu().detach()

					# get labels from vocab
					_A :				Any				=  ([''''''] + list(self.processor.tokenizer.get_vocab().keys()			))[
					    :-1
					]  # logits don't align with the tokenizer's vocab

					_A :				Union[str, Any]				=  {c: i for i, c in enumerate(UpperCAmelCase__			)}
					_A :				List[str]				=  []
					for c in transcript:
									if c in dictionary:
													tokens.append(dictionary[c]			)

					def get_trellis(UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__=0			):
									_A :				Dict				=  emission.size(0			)
									_A :				List[Any]				=  len(UpperCAmelCase__			)

									# 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.
									_A :				Tuple				=  torch.full((num_frames + 1, num_tokens + 1)   ,     -float('''inf'''			)			)
									_A :				Any				=  0
									for t in range(UpperCAmelCase__			):
													_A :				str				=  torch.maximum(
													    # Score for staying at the same token
													    trellis[t, 1:] + emission[t, blank_id]   ,     trellis[t, :-1] + emission[t, tokens]   ,     )
									return trellis

					_A :				Tuple				=  get_trellis(UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)



					@dataclass
					class  lowerCamelCase__				:


						"""simple docstring"""





						__magic_name__							   =		42
						__magic_name__							   =		42
						__magic_name__							   =		42

					def backtrack(UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__=0			):
									_A :				Optional[int]				=  trellis.size(1			) - 1
									_A :				List[str]				=  torch.argmax(trellis[:, j]			).item()

									_A :				int				=  []
									for t in range(UpperCAmelCase__   ,     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.
													_A :				Any				=  trellis[t - 1, j] + emission[t - 1, blank_id]
													# Score for token changing from C-1 at T-1 to J at T.
													_A :				Optional[Any]				=  trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]

													# 2. Store the path with frame-wise probability.
													_A :				str				=  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   ,     UpperCAmelCase__			)			)

													# 3. Update the token
													if changed > stayed:
																	j -= 1
																	if j == 0:
																					break
									else:
													raise ValueError('''Failed to align'''			)
									return path[::-1]

					_A :				int				=  backtrack(UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)



					@dataclass
					class  lowerCamelCase__				:


						"""simple docstring"""





						__magic_name__							   =		42
						__magic_name__							   =		42
						__magic_name__							   =		42
						__magic_name__							   =		42




						def __repr__(					self			)       ->				str:
										return F"""{self.label}\t{self.score:4.2f}\t{self.start*2_0:5d}\t{self.end*2_0:5d}"""




						@property
						def 				_lowerCamelCase				(					self			)       ->				List[str]:
										return self.end - self.start

					def merge_repeats(UpperCAmelCase__			):
									_A	,       _A :				Any				=  0, 0
									_A :				List[Any]				=  []
									while ia < len(UpperCAmelCase__			):
													while ia < len(UpperCAmelCase__			) and path[ia].token_index == path[ia].token_index:
																	ia += 1
													_A :				Dict				=  sum(path[k].score for k in range(UpperCAmelCase__   ,     UpperCAmelCase__			)			) / (ia - ia)
													segments.append(
													    Segment(
													        transcript[path[ia].token_index]   ,     path[ia].time_index   ,     path[ia - 1].time_index + 1   ,     UpperCAmelCase__   ,     )			)
													_A :				Dict				=  ia
									return segments

					_A :				Optional[Any]				=  merge_repeats(UpperCAmelCase__			)
					with open(item['''out_path''']   ,     '''w'''			) as out_align:
									for seg in segments:
													out_align.write(str(UpperCAmelCase__			) + '''\n'''			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)       ->				Tuple:
					if not os.path.exists(UpperCAmelCase__			):
									os.makedirs(UpperCAmelCase__			)

					# load text file
					_A :				List[str]				=  open(UpperCAmelCase__   ,     encoding='''utf8'''			).readlines()

					_A :				Optional[int]				=  []
					for line in lines:
									if len(line.strip().split('''\t'''			)			) != 2:
													print('''Script must be in format: 00001  this is my sentence'''			)
													exit()

									_A	,       _A :				Dict				=  line.strip().split('''\t'''			)
									_A :				Dict				=  os.path.join(UpperCAmelCase__   ,     wav_name + '''.wav'''			)
									_A :				Tuple				=  os.path.join(UpperCAmelCase__   ,     wav_name + '''.txt'''			)

									items.append({'''sent''': sentence, '''wav_path''': wav_path, '''out_path''': out_path}			)
					print('''Number of samples found in script file'''   ,     len(UpperCAmelCase__			)			)

					for item in tqdm(UpperCAmelCase__			):
									self.align_single_sample(UpperCAmelCase__			)





def      lowercase  (  ):



				"""simple docstring"""


				_A :				str				=  argparse.ArgumentParser()

				parser.add_argument(
				    '''--model_name'''   , type=lowerCAmelCase   , default='''arijitx/wav2vec2-xls-r-300m-bengali'''   , help='''wav2vec model name''')
				parser.add_argument('''--wav_dir'''   , type=lowerCAmelCase   , default='''./wavs'''   , help='''directory containing wavs''')
				parser.add_argument('''--text_file'''   , type=lowerCAmelCase   , default='''script.txt'''   , help='''file containing text''')
				parser.add_argument('''--input_wavs_sr'''   , type=lowerCAmelCase   , default=1_6000   , help='''sampling rate of input audios''')
				parser.add_argument(
				    '''--output_dir'''   , type=lowerCAmelCase   , default='''./out_alignment'''   , help='''output directory containing the alignment files''')
				parser.add_argument('''--cuda'''   , action='''store_true''')

				_A :				Any				=  parser.parse_args()

				_A :				Any				=  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()







'''simple docstring'''






import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    ElectraForCausalLM,
	    ElectraForMaskedLM,
	    ElectraForMultipleChoice,
	    ElectraForPreTraining,
	    ElectraForQuestionAnswering,
	    ElectraForSequenceClassification,
	    ElectraForTokenClassification,
	    ElectraModel,
	)
	from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST



class  lowerCamelCase__				:


	"""simple docstring"""





	def __init__(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__=1_3   ,     UpperCAmelCase__=7   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=9_9   ,     UpperCAmelCase__=3_2   ,     UpperCAmelCase__=5   ,     UpperCAmelCase__=4   ,     UpperCAmelCase__=3_7   ,     UpperCAmelCase__="gelu"   ,     UpperCAmelCase__=0.1   ,     UpperCAmelCase__=0.1   ,     UpperCAmelCase__=5_1_2   ,     UpperCAmelCase__=1_6   ,     UpperCAmelCase__=2   ,     UpperCAmelCase__=0.0_2   ,     UpperCAmelCase__=3   ,     UpperCAmelCase__=4   ,     UpperCAmelCase__=None   ,     )       ->				Optional[Any]:
					_A :				Optional[int]				=  parent
					_A :				List[Any]				=  batch_size
					_A :				Optional[int]				=  seq_length
					_A :				Dict				=  is_training
					_A :				Union[str, Any]				=  use_input_mask
					_A :				Tuple				=  use_token_type_ids
					_A :				str				=  use_labels
					_A :				List[str]				=  vocab_size
					_A :				Tuple				=  hidden_size
					_A :				Tuple				=  num_hidden_layers
					_A :				List[str]				=  num_attention_heads
					_A :				Optional[Any]				=  intermediate_size
					_A :				str				=  hidden_act
					_A :				Optional[int]				=  hidden_dropout_prob
					_A :				int				=  attention_probs_dropout_prob
					_A :				Optional[Any]				=  max_position_embeddings
					_A :				List[Any]				=  type_vocab_size
					_A :				Tuple				=  type_sequence_label_size
					_A :				Any				=  initializer_range
					_A :				Any				=  num_labels
					_A :				str				=  num_choices
					_A :				List[str]				=  scope




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				List[Any]				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size			)

					_A :				Union[str, Any]				=  None
					if self.use_input_mask:
									_A :				Tuple				=  random_attention_mask([self.batch_size, self.seq_length]			)

					_A :				List[str]				=  None
					if self.use_token_type_ids:
									_A :				List[Any]				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.type_vocab_size			)

					_A :				Tuple				=  None
					_A :				Dict				=  None
					_A :				int				=  None
					if self.use_labels:
									_A :				Any				=  ids_tensor([self.batch_size]   ,     self.type_sequence_label_size			)
									_A :				Dict				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.num_labels			)
									_A :				int				=  ids_tensor([self.batch_size]   ,     self.num_choices			)
									_A :				List[str]				=  ids_tensor([self.batch_size, self.seq_length]   ,     1			)

					_A :				List[Any]				=  self.get_config()

					return (
					    config,
					    input_ids,
					    token_type_ids,
					    input_mask,
					    sequence_labels,
					    token_labels,
					    choice_labels,
					    fake_token_labels,
					)




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					return ElectraConfig(
					    vocab_size=self.vocab_size   ,     hidden_size=self.hidden_size   ,     num_hidden_layers=self.num_hidden_layers   ,     num_attention_heads=self.num_attention_heads   ,     intermediate_size=self.intermediate_size   ,     hidden_act=self.hidden_act   ,     hidden_dropout_prob=self.hidden_dropout_prob   ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,     max_position_embeddings=self.max_position_embeddings   ,     type_vocab_size=self.type_vocab_size   ,     is_decoder=UpperCAmelCase__   ,     initializer_range=self.initializer_range   ,     )




	def 				_lowerCamelCase				(					self			)       ->				int:
					(
					    (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       
					) :				List[str]				=  self.prepare_config_and_inputs()

					_A :				Tuple				=  True
					_A :				List[str]				=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
					_A :				str				=  ids_tensor([self.batch_size, self.seq_length]   ,     vocab_size=2			)

					return (
					    config,
					    input_ids,
					    token_type_ids,
					    input_mask,
					    sequence_labels,
					    token_labels,
					    choice_labels,
					    encoder_hidden_states,
					    encoder_attention_mask,
					)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Tuple:
					_A :				List[str]				=  ElectraModel(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Tuple				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					_A :				List[Any]				=  model(UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					_A :				Tuple				=  model(UpperCAmelCase__			)
					self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				Dict				=  True
					_A :				Any				=  ElectraModel(UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				List[str]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     encoder_hidden_states=UpperCAmelCase__   ,     encoder_attention_mask=UpperCAmelCase__   ,     )
					_A :				Optional[int]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     encoder_hidden_states=UpperCAmelCase__   ,     )
					_A :				List[Any]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				Optional[int]				=  ElectraForMaskedLM(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				List[str]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.vocab_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Union[str, Any]:
					_A :				List[Any]				=  ElectraForCausalLM(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Optional[int]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.vocab_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				str				=  self.num_labels
					_A :				str				=  ElectraForTokenClassification(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				str				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.num_labels)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[Any]:
					_A :				Optional[int]				=  self.num_labels
					_A :				Optional[int]				=  ElectraForPreTraining(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Dict				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Tuple:
					_A :				Optional[Any]				=  self.num_labels
					_A :				Optional[Any]				=  ElectraForSequenceClassification(UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Any				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_labels)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[Any]:
					_A :				Optional[int]				=  ElectraForQuestionAnswering(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Union[str, Any]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     start_positions=UpperCAmelCase__   ,     end_positions=UpperCAmelCase__   ,     )
					self.parent.assertEqual(result.start_logits.shape   ,     (self.batch_size, self.seq_length)			)
					self.parent.assertEqual(result.end_logits.shape   ,     (self.batch_size, self.seq_length)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Union[str, Any]:
					_A :				Union[str, Any]				=  self.num_choices
					_A :				Union[str, Any]				=  ElectraForMultipleChoice(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Dict				=  input_ids.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				str				=  token_type_ids.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				Dict				=  input_mask.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				str				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__   ,     )
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_choices)			)




	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					_A :				Optional[Any]				=  self.prepare_config_and_inputs()
					(
					    (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       
					) :				int				=  config_and_inputs
					_A :				Dict				=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
					return config, inputs_dict



@require_torch
class  lowerCamelCase__				(						snake_case_  ,							snake_case_  ,							unittest.TestCase						):


	"""simple docstring"""





	__magic_name__							   =		(
	    (
	        ElectraModel,
	        ElectraForPreTraining,
	        ElectraForMaskedLM,
	        ElectraForCausalLM,
	        ElectraForMultipleChoice,
	        ElectraForTokenClassification,
	        ElectraForSequenceClassification,
	        ElectraForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	__magic_name__							   =		(
	    {
	        """feature-extraction""": ElectraModel,
	        """fill-mask""": ElectraForMaskedLM,
	        """question-answering""": ElectraForQuestionAnswering,
	        """text-classification""": ElectraForSequenceClassification,
	        """text-generation""": ElectraForCausalLM,
	        """token-classification""": ElectraForTokenClassification,
	        """zero-shot""": ElectraForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	__magic_name__							   =		True




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__=False			)       ->				Optional[Any]:
					_A :				List[str]				=  super()._prepare_for_class(UpperCAmelCase__   ,     UpperCAmelCase__   ,     return_labels=UpperCAmelCase__			)

					if return_labels:
									if model_class in get_values(UpperCAmelCase__			):
													_A :				Tuple				=  torch.zeros(
													    (self.model_tester.batch_size, self.model_tester.seq_length)   ,     dtype=torch.long   ,     device=UpperCAmelCase__			)
					return inputs_dict




	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Optional[Any]				=  ElectraModelTester(self			)
					_A :				List[str]				=  ConfigTester(self   ,     config_class=UpperCAmelCase__   ,     hidden_size=3_7			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					self.config_tester.run_common_tests()




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				Union[str, Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_model(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_electra_model_as_decoder(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				int				=  self.model_tester.prepare_config_and_inputs()
					for type in ["absolute", "relative_key", "relative_key_query"]:
									_A :				List[str]				=  type
									self.model_tester.create_and_check_electra_model(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_masked_lm(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Any:
					_A :				Dict				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_token_classification(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Union[str, Any]:
					_A :				List[str]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_pretraining(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				Any				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_sequence_classification(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_question_answering(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_multiple_choice(*UpperCAmelCase__			)




	@slow
	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_A :				List[str]				=  ElectraModel.from_pretrained(UpperCAmelCase__			)
									self.assertIsNotNone(UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				str				=  self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_electra_for_causal_lm(*UpperCAmelCase__			)



@require_torch
class  lowerCamelCase__				(						unittest.TestCase						):


	"""simple docstring"""





	@slow
	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Any				=  ElectraModel.from_pretrained('''google/electra-small-discriminator'''			)
					_A :				Union[str, Any]				=  torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]			)
					_A :				Union[str, Any]				=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
					_A :				Optional[int]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__			)[0]
					_A :				Tuple				=  torch.Size((1, 1_1, 2_5_6)			)
					self.assertEqual(output.shape   ,     UpperCAmelCase__			)
					_A :				Optional[Any]				=  torch.tensor(
					    [[[0.4_4_7_1, 0.6_8_2_1, -0.3_2_6_5], [0.4_6_2_7, 0.5_2_5_5, -0.3_6_6_8], [0.4_5_3_2, 0.3_3_1_3, -0.4_3_4_4]]]			)

					self.assertTrue(torch.allclose(output[:, 1:4, 1:4]   ,     UpperCAmelCase__   ,     atol=1e-4			)			)
