






import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
  import torch
  from torch import nn

  from transformers import YolosForObjectDetection, YolosModel
  from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor


class       __lowercase      :







     '''simple docstring'''




     def __init__(				self       :			Any ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Optional[int]=13 ,	UpperCamelCase_       :			Any=[30, 30] ,	UpperCamelCase_       :			Optional[Any]=2 ,	UpperCamelCase_       :			List[Any]=3 ,	UpperCamelCase_       :			Union[str, Any]=True ,	UpperCamelCase_       :			str=True ,	UpperCamelCase_       :			int=32 ,	UpperCamelCase_       :			Optional[Any]=5 ,	UpperCamelCase_       :			List[str]=4 ,	UpperCamelCase_       :			str=37 ,	UpperCamelCase_       :			Any="gelu" ,	UpperCamelCase_       :			str=0.1 ,	UpperCamelCase_       :			Optional[int]=0.1 ,	UpperCamelCase_       :			Dict=10 ,	UpperCamelCase_       :			str=0.02 ,	UpperCamelCase_       :			Optional[Any]=3 ,	UpperCamelCase_       :			List[str]=None ,	UpperCamelCase_       :			Dict=8 ,	UpperCamelCase_       :			List[str]=10 ,	):


         """simple docstring"""



         __A         =	parent
         __A         =	batch_size
         __A         =	image_size
         __A         =	patch_size
         __A         =	num_channels
         __A         =	is_training
         __A         =	use_labels
         __A         =	hidden_size
         __A         =	num_hidden_layers
         __A         =	num_attention_heads
         __A         =	intermediate_size
         __A         =	hidden_act
         __A         =	hidden_dropout_prob
         __A         =	attention_probs_dropout_prob
         __A         =	type_sequence_label_size
         __A         =	initializer_range
         __A         =	num_labels
         __A         =	scope
         __A         =	n_targets
         __A         =	num_detection_tokens
         # we set the expected sequence length (which is used in several tests)
         # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
         __A         =	(image_size[1] // patch_size) * (image_size[0] // patch_size)
         __A         =	num_patches + 1 + self.num_detection_tokens




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]				)

         __A         =	None
         if self.use_labels:
             # labels is a list of Dict (each Dict being the labels for a given example in the batch)
             __A         =	[]
             for i in range(self.batch_size				):
                 __A         =	{}
                 __A         =	torch.randint(
                     high=self.num_labels ,	size=(self.n_targets,) ,	device=UpperCamelCase_				)
                 __A         =	torch.rand(self.n_targets ,	4 ,	device=UpperCamelCase_				)
                 labels.append(UpperCamelCase_				)

         __A         =	self.get_config()

         return config, pixel_values, labels




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         return YolosConfig(
             image_size=self.image_size ,	patch_size=self.patch_size ,	num_channels=self.num_channels ,	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 ,	is_decoder=UpperCamelCase_ ,	initializer_range=self.initializer_range ,	num_detection_tokens=self.num_detection_tokens ,	num_labels=self.num_labels ,	)




     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Optional[Any]				):


         """simple docstring"""



         __A         =	YolosModel(config=UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.eval()
         __A         =	model(UpperCamelCase_				)
         self.parent.assertEqual(
             result.last_hidden_state.shape ,	(self.batch_size, self.expected_seq_len, self.hidden_size)				)




     def 		lowerCAmelCase_     (				self       :			Any ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Any				):


         """simple docstring"""



         __A         =	YolosForObjectDetection(UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.eval()

         __A         =	model(pixel_values=UpperCamelCase_				)
         __A         =	model(UpperCamelCase_				)

         self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_detection_tokens, self.num_labels + 1)				)
         self.parent.assertEqual(result.pred_boxes.shape ,	(self.batch_size, self.num_detection_tokens, 4)				)

         __A         =	model(pixel_values=UpperCamelCase_ ,	labels=UpperCamelCase_				)

         self.parent.assertEqual(result.loss.shape ,	()				)
         self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_detection_tokens, self.num_labels + 1)				)
         self.parent.assertEqual(result.pred_boxes.shape ,	(self.batch_size, self.num_detection_tokens, 4)				)





     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A         =	self.prepare_config_and_inputs()
         __A    ,	__A    ,	__A         =	config_and_inputs
         __A         =	{"""pixel_values""": pixel_values}
         return config, inputs_dict


@require_torch
class       __lowercase      (	lowercase_			,      lowercase_			,      unittest.TestCase  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				(YolosModel, YolosForObjectDetection) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE   						=				(
         {"feature-extraction": YolosModel, "object-detection": YolosForObjectDetection} if is_torch_available() else {}
     )

     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False




     def 		lowerCAmelCase_     (				self       :			Dict ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Any=False				):


         """simple docstring"""



         __A         =	super()._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_ ,	return_labels=UpperCamelCase_				)

         if return_labels:
             if model_class.__name__ == "YolosForObjectDetection":
                 __A         =	[]
                 for i in range(self.model_tester.batch_size				):
                     __A         =	{}
                     __A         =	torch.ones(
                         size=(self.model_tester.n_targets,) ,	device=UpperCamelCase_ ,	dtype=torch.long				)
                     __A         =	torch.ones(
                         self.model_tester.n_targets ,	4 ,	device=UpperCamelCase_ ,	dtype=torch.float				)
                     labels.append(UpperCamelCase_				)
                 __A         =	labels

         return inputs_dict




     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         __A         =	YolosModelTester(self				)
         __A         =	ConfigTester(self ,	config_class=UpperCamelCase_ ,	has_text_modality=UpperCamelCase_ ,	hidden_size=37				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         self.config_tester.run_common_tests()




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         pass




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             __A         =	model_class(UpperCamelCase_				)
             self.assertIsInstance(model.get_input_embeddings() ,	(nn.Module)				)
             __A         =	model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(UpperCamelCase_ ,	nn.Linear				)				)




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             __A         =	model_class(UpperCamelCase_				)
             __A         =	inspect.signature(model.forward				)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             __A         =	[*signature.parameters.keys()]

             __A         =	["""pixel_values"""]
             self.assertListEqual(arg_names[:1] ,	UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         __A         =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()
         __A         =	True

         # in YOLOS, the seq_len is different
         __A         =	self.model_tester.expected_seq_len
         for model_class in self.all_model_classes:
             __A         =	True
             __A         =	False
             __A         =	True
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()
             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)
             __A         =	outputs.attentions
             self.assertEqual(len(UpperCamelCase_				) ,	self.model_tester.num_hidden_layers				)

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             __A         =	True
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()
             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)
             __A         =	outputs.attentions
             self.assertEqual(len(UpperCamelCase_				) ,	self.model_tester.num_hidden_layers				)

             self.assertListEqual(
                 list(attentions[0].shape[-3:]				) ,	[self.model_tester.num_attention_heads, seq_len, seq_len] ,	)
             __A         =	len(UpperCamelCase_				)

             # Check attention is always last and order is fine
             __A         =	True
             __A         =	True
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()
             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)

             __A         =	1
             self.assertEqual(out_len + added_hidden_states ,	len(UpperCamelCase_				)				)

             __A         =	outputs.attentions

             self.assertEqual(len(UpperCamelCase_				) ,	self.model_tester.num_hidden_layers				)
             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]				) ,	[self.model_tester.num_attention_heads, seq_len, seq_len] ,	)




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         def check_hidden_states_output(UpperCamelCase_       :			int ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Tuple				):
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()

             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)

             __A         =	outputs.hidden_states

             __A         =	getattr(
                 self.model_tester ,	"""expected_num_hidden_layers""" ,	self.model_tester.num_hidden_layers + 1				)
             self.assertEqual(len(UpperCamelCase_				) ,	UpperCamelCase_				)

             # YOLOS has a different seq_length
             __A         =	self.model_tester.expected_seq_len

             self.assertListEqual(
                 list(hidden_states[0].shape[-2:]				) ,	[seq_length, self.model_tester.hidden_size] ,	)

         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             __A         =	True
             check_hidden_states_output(UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_				)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             __A         =	True

             check_hidden_states_output(UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



         __A         =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_object_detection(*UpperCamelCase_				)





     @slow
     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A         =	YolosModel.from_pretrained(UpperCamelCase_				)
             self.assertIsNotNone(UpperCamelCase_				)



def 			_SCREAMING_SNAKE_CASE						( )						->      str:





    """simple docstring"""

    __A         =	Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
    return image






@require_torch
@require_vision
class       __lowercase      (	unittest.TestCase  ):







     '''simple docstring'''




     @cached_property
     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         return AutoImageProcessor.from_pretrained("""hustvl/yolos-small"""				) if is_vision_available() else None





     @slow
     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         __A         =	YolosForObjectDetection.from_pretrained("""hustvl/yolos-small"""				).to(UpperCamelCase_				)

         __A         =	self.default_image_processor
         __A         =	prepare_img()
         __A         =	image_processor(images=UpperCamelCase_ ,	return_tensors="""pt"""				).to(UpperCamelCase_				)

         # forward pass
         with torch.no_grad():
             __A         =	model(inputs.pixel_values				)

         # verify outputs
         __A         =	torch.Size((1, 100, 92)				)
         self.assertEqual(outputs.logits.shape ,	UpperCamelCase_				)

         __A         =	torch.tensor(
             [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]] ,	device=UpperCamelCase_ ,	)
         __A         =	torch.tensor(
             [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]] ,	device=UpperCamelCase_				)
         self.assertTrue(torch.allclose(outputs.logits[0, :3, :3] ,	UpperCamelCase_ ,	atol=1e-4				)				)
         self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3] ,	UpperCamelCase_ ,	atol=1e-4				)				)

         # verify postprocessing
         __A         =	image_processor.post_process_object_detection(
             UpperCamelCase_ ,	threshold=0.3 ,	target_sizes=[image.size[::-1]]				)[0]
         __A         =	torch.tensor([0.9994, 0.9790, 0.9964, 0.9972, 0.9861]				).to(UpperCamelCase_				)
         __A         =	[75, 75, 17, 63, 17]
         __A         =	torch.tensor([335.0609, 79.3848, 375.4216, 187.2495]				).to(UpperCamelCase_				)

         self.assertEqual(len(results["""scores"""]				) ,	5				)
         self.assertTrue(torch.allclose(results["""scores"""] ,	UpperCamelCase_ ,	atol=1e-4				)				)
         self.assertSequenceEqual(results["""labels"""].tolist() ,	UpperCamelCase_				)
         self.assertTrue(torch.allclose(results["""boxes"""][0, :] ,	UpperCamelCase_				)				)












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

import datasets
import numpy as np
import torch
import torchaudio
from packaging import version
from torch import nn

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaProcessor,
    is_apex_available,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


if is_apex_available():
  from apex import amp


if version.parse(version.parse(torch.__version__).base_version) >= version.parse("1.6"):
  __a      :       Optional[Any]				  =   True
  from torch.cuda.amp import autocast

__a      :       Any				  =   logging.getLogger(__name__)



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			List[str]=None			,				__lowercase   :			Tuple=None    )						->      Dict:





    """simple docstring"""

    return field(default_factory=lambda: default			,				metadata=__lowercase    )


@dataclass
class       __lowercase      :







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				field(
         metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}  )
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}			,      )
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={"help": "Whether to freeze the feature extractor layers of the model."}  )
     SCREAMING_SNAKE_CASE   						=				field(
         default=0.1			,      metadata={"help": "The dropout ratio for the attention probabilities."}  )
     SCREAMING_SNAKE_CASE   						=				field(
         default=0.1			,      metadata={"help": "The dropout ratio for activations inside the fully connected layer."}  )
     SCREAMING_SNAKE_CASE   						=				field(
         default=0.1			,      metadata={
             "help": "The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler."
         }			,      )
     SCREAMING_SNAKE_CASE   						=				field(
         default=0.1			,      metadata={"help": "The dropout probabilitiy for all 1D convolutional layers in feature extractor."}			,      )
     SCREAMING_SNAKE_CASE   						=				field(
         default=0.05			,      metadata={
             "help": (
                 "Propability of each feature vector along the time axis to be chosen as the start of the vector"
                 "span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature"
                 "vectors will be masked along the time axis. This is only relevant if ``apply_spec_augment is True``."
             )
         }			,      )
     SCREAMING_SNAKE_CASE   						=				field(default=0.0			,      metadata={"help": "The LayerDrop probability."}  )


@dataclass
class       __lowercase      :







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}  )
     SCREAMING_SNAKE_CASE   						=				field(
         default="train+validation"			,      metadata={
             "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
         }			,      )
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={"help": "Overwrite the cached preprocessed datasets or not."}  )
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={"help": "The number of processes to use for the preprocessing."}			,      )
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={
             "help": (
                 "For debugging purposes or quicker training, truncate the number of training examples to this "
                 "value if set."
             )
         }			,      )
     SCREAMING_SNAKE_CASE   						=				field(
         default=lowercase_			,      metadata={
             "help": (
                 "For debugging purposes or quicker training, truncate the number of validation examples to this "
                 "value if set."
             )
         }			,      )
     SCREAMING_SNAKE_CASE   						=				list_field(
         default=[",", "?", ".", "!", "-", ";", ":", "\"\"", "%", "'", "\"", "�"]			,      metadata={"help": "A list of characters to remove from the transcripts."}			,      )


@dataclass
class       __lowercase      :







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				42
     SCREAMING_SNAKE_CASE   						=				True
     SCREAMING_SNAKE_CASE   						=				None
     SCREAMING_SNAKE_CASE   						=				None
     SCREAMING_SNAKE_CASE   						=				None
     SCREAMING_SNAKE_CASE   						=				None




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


         """simple docstring"""



         __A         =	[{"""input_values""": feature["""input_values"""]} for feature in features]
         __A         =	[{"""input_ids""": feature["""labels"""]} for feature in features]

         __A         =	self.processor.pad(
             UpperCamelCase_ ,	padding=self.padding ,	max_length=self.max_length ,	pad_to_multiple_of=self.pad_to_multiple_of ,	return_tensors="""pt""" ,	)
         __A         =	self.processor.pad(
             labels=UpperCamelCase_ ,	padding=self.padding ,	max_length=self.max_length_labels ,	pad_to_multiple_of=self.pad_to_multiple_of_labels ,	return_tensors="""pt""" ,	)

         # replace padding with -100 to ignore loss correctly
         __A         =	labels_batch["""input_ids"""].masked_fill(labels_batch.attention_mask.ne(1				) ,	-100				)

         __A         =	labels

         return batch


class       __lowercase      (	lowercase_  ):







     '''simple docstring'''




     def 		lowerCAmelCase_     (				self       :			List[str] ,	UpperCamelCase_       :			nn.Module ,	UpperCamelCase_       :			Dict[str, Union[torch.Tensor, Any]]				):


         """simple docstring"""



         model.train()
         __A         =	self._prepare_inputs(UpperCamelCase_				)

         if self.use_amp:
             with autocast():
                 __A         =	self.compute_loss(UpperCamelCase_ ,	UpperCamelCase_				)
         else:
             __A         =	self.compute_loss(UpperCamelCase_ ,	UpperCamelCase_				)

         if self.args.n_gpu > 1:
             if model.module.config.ctc_loss_reduction == "mean":
                 __A         =	loss.mean()
             elif model.module.config.ctc_loss_reduction == "sum":
                 __A         =	loss.sum() / (inputs["""labels"""] >= 0).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:
             __A         =	loss / self.args.gradient_accumulation_steps

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

         return loss.detach()



def 			_SCREAMING_SNAKE_CASE						( )						->      Optional[int]:





    """simple docstring"""

    __A         =	HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)    )
    if len(sys.argv    ) == 2 and sys.argv[1].endswith(""".json"""    ):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        __A    ,	__A    ,	__A         =	parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]    )    )
    else:
        __A    ,	__A    ,	__A         =	parser.parse_args_into_dataclasses()

    # Detecting last checkpoint.
    __A         =	None
    if os.path.isdir(training_args.output_dir    ) and training_args.do_train and not training_args.overwrite_output_dir:
        __A         =	get_last_checkpoint(training_args.output_dir    )
        if last_checkpoint is None and len(os.listdir(training_args.output_dir    )    ) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                """Use --overwrite_output_dir to overcome."""    )
        elif last_checkpoint is not None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                """the `--output_dir` or add `--overwrite_output_dir` to train from scratch."""    )

    # Setup logging
    logging.basicConfig(
        format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s"""			,				datefmt="""%m/%d/%Y %H:%M:%S"""			,				handlers=[logging.StreamHandler(sys.stdout    )]			,				)
    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank    ) else logging.WARN    )

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1    )}, 16-bits training: {training_args.fpaa}"    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank    ):
        transformers.utils.logging.set_verbosity_info()
    logger.info("""Training/evaluation parameters %s"""			,				__lowercase    )

    # Set seed before initializing model.
    set_seed(training_args.seed    )

    # Get the datasets:
    __A         =	datasets.load_dataset(
        """common_voice"""			,				data_args.dataset_config_name			,				split=data_args.train_split_name    )
    __A         =	datasets.load_dataset("""common_voice"""			,				data_args.dataset_config_name			,				split="""test"""    )

    # Create and save tokenizer
    __A         =	f"[{''.join(data_args.chars_to_ignore    )}]"

    def remove_special_characters(__lowercase   :			Any    ):
        __A         =	re.sub(__lowercase			,				""""""			,				batch["""sentence"""]    ).lower() + """ """
        return batch

    __A         =	train_dataset.map(__lowercase			,				remove_columns=["""sentence"""]    )
    __A         =	eval_dataset.map(__lowercase			,				remove_columns=["""sentence"""]    )

    def extract_all_chars(__lowercase   :			Tuple    ):
        __A         =	""" """.join(batch["""text"""]    )
        __A         =	list(set(__lowercase    )    )
        return {"vocab": [vocab], "all_text": [all_text]}

    __A         =	train_dataset.map(
        __lowercase			,				batched=__lowercase			,				batch_size=-1			,				keep_in_memory=__lowercase			,				remove_columns=train_dataset.column_names			,				)
    __A         =	train_dataset.map(
        __lowercase			,				batched=__lowercase			,				batch_size=-1			,				keep_in_memory=__lowercase			,				remove_columns=eval_dataset.column_names			,				)

    __A         =	list(set(vocab_train["""vocab"""][0]    ) | set(vocab_test["""vocab"""][0]    )    )
    __A         =	{v: k for k, v in enumerate(__lowercase    )}
    __A         =	vocab_dict[""" """]
    del vocab_dict[" "]
    __A         =	len(__lowercase    )
    __A         =	len(__lowercase    )

    with open("""vocab.json"""			,				"""w"""    ) as vocab_file:
        json.dump(__lowercase			,				__lowercase    )

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    __A         =	WavaVecaCTCTokenizer(
        """vocab.json"""			,				unk_token="""[UNK]"""			,				pad_token="""[PAD]"""			,				word_delimiter_token="""|"""			,				)
    __A         =	WavaVecaFeatureExtractor(
        feature_size=1			,				sampling_rate=1_6_0_0_0			,				padding_value=0.0			,				do_normalize=__lowercase			,				return_attention_mask=__lowercase    )
    __A         =	WavaVecaProcessor(feature_extractor=__lowercase			,				tokenizer=__lowercase    )
    __A         =	WavaVecaForCTC.from_pretrained(
        model_args.model_name_or_path			,				cache_dir=model_args.cache_dir			,				activation_dropout=model_args.activation_dropout			,				attention_dropout=model_args.attention_dropout			,				hidden_dropout=model_args.hidden_dropout			,				feat_proj_dropout=model_args.feat_proj_dropout			,				mask_time_prob=model_args.mask_time_prob			,				gradient_checkpointing=training_args.gradient_checkpointing			,				layerdrop=model_args.layerdrop			,				ctc_loss_reduction="""mean"""			,				pad_token_id=processor.tokenizer.pad_token_id			,				vocab_size=len(processor.tokenizer    )			,				)

    if data_args.max_train_samples is not None:
        __A         =	min(len(__lowercase    )			,				data_args.max_train_samples    )
        __A         =	train_dataset.select(range(__lowercase    )    )

    if data_args.max_val_samples is not None:
        __A         =	eval_dataset.select(range(data_args.max_val_samples    )    )

    __A         =	torchaudio.transforms.Resample(4_8_0_0_0			,				1_6_0_0_0    )

    # Preprocessing the datasets.
    # We need to read the aduio files as arrays and tokenize the targets.
    def speech_file_to_array_fn(__lowercase   :			Tuple    ):
        __A    ,	__A         =	torchaudio.load(batch["""path"""]    )
        __A         =	resampler(__lowercase    ).squeeze().numpy()
        __A         =	1_6_0_0_0
        __A         =	batch["""text"""]
        return batch

    __A         =	train_dataset.map(
        __lowercase			,				remove_columns=train_dataset.column_names			,				num_proc=data_args.preprocessing_num_workers			,				)
    __A         =	eval_dataset.map(
        __lowercase			,				remove_columns=eval_dataset.column_names			,				num_proc=data_args.preprocessing_num_workers			,				)

    def prepare_dataset(__lowercase   :			List[Any]    ):
        # check that all files have the correct sampling rate
        assert (
            len(set(batch["""sampling_rate"""]    )    ) == 1
        ), f"Make sure all inputs have the same sampling rate of {processor.feature_extractor.sampling_rate}."

        __A         =	processor(
            audio=batch["""speech"""]			,				text=batch["""target_text"""]			,				sampling_rate=batch["""sampling_rate"""][0]    )
        batch.update(__lowercase    )
        return batch

    __A         =	train_dataset.map(
        __lowercase			,				remove_columns=train_dataset.column_names			,				batch_size=training_args.per_device_train_batch_size			,				batched=__lowercase			,				num_proc=data_args.preprocessing_num_workers			,				)
    __A         =	eval_dataset.map(
        __lowercase			,				remove_columns=eval_dataset.column_names			,				batch_size=training_args.per_device_train_batch_size			,				batched=__lowercase			,				num_proc=data_args.preprocessing_num_workers			,				)

    # Metric
    __A         =	datasets.load_metric("""wer"""    )

    def compute_metrics(__lowercase   :			Union[str, Any]    ):
        __A         =	pred.predictions
        __A         =	np.argmax(__lowercase			,				axis=-1    )

        __A         =	processor.tokenizer.pad_token_id

        __A         =	processor.batch_decode(__lowercase    )
        # we do not want to group tokens when computing the metrics
        __A         =	processor.batch_decode(pred.label_ids			,				group_tokens=__lowercase    )

        __A         =	wer_metric.compute(predictions=__lowercase			,				references=__lowercase    )

        return {"wer": wer}

    if model_args.freeze_feature_extractor:
        model.freeze_feature_extractor()

    # Data collator
    __A         =	DataCollatorCTCWithPadding(processor=__lowercase			,				padding=__lowercase    )

    # Initialize our Trainer
    __A         =	CTCTrainer(
        model=__lowercase			,				data_collator=__lowercase			,				args=__lowercase			,				compute_metrics=__lowercase			,				train_dataset=train_dataset if training_args.do_train else None			,				eval_dataset=eval_dataset if training_args.do_eval else None			,				tokenizer=processor.feature_extractor			,				)

    # Training
    if training_args.do_train:
        if last_checkpoint is not None:
            __A         =	last_checkpoint
        elif os.path.isdir(model_args.model_name_or_path    ):
            __A         =	model_args.model_name_or_path
        else:
            __A         =	None

        # Save the feature_extractor and the tokenizer
        if is_main_process(training_args.local_rank    ):
            processor.save_pretrained(training_args.output_dir    )

        __A         =	trainer.train(resume_from_checkpoint=__lowercase    )
        trainer.save_model()

        __A         =	train_result.metrics
        __A         =	(
            data_args.max_train_samples if data_args.max_train_samples is not None else len(__lowercase    )
        )
        __A         =	min(__lowercase			,				len(__lowercase    )    )

        trainer.log_metrics("""train"""			,				__lowercase    )
        trainer.save_metrics("""train"""			,				__lowercase    )
        trainer.save_state()

    # Evaluation
    __A         =	{}
    if training_args.do_eval:
        logger.info("""*** Evaluate ***"""    )
        __A         =	trainer.evaluate()
        __A         =	data_args.max_val_samples if data_args.max_val_samples is not None else len(__lowercase    )
        __A         =	min(__lowercase			,				len(__lowercase    )    )

        trainer.log_metrics("""eval"""			,				__lowercase    )
        trainer.save_metrics("""eval"""			,				__lowercase    )

    return results


if __name__ == "__main__":
  main()




