import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
	from huggingface_hub import snapshot_download
	from pyctcdecode import BeamSearchDecoderCTC

	from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
	from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
	from transformers import WavaVecaForCTC





@require_pyctcdecode
class UpperCAmelCase		(   unittest.TestCase    ):


				'''simple docstring'''






				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Any      = '| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
						A_						:							Dict      = dict(zip(lowercase						,      range(len(lowercase   )   )   )   )

						A_						:							List[Any]      = {
						    'unk_token': '<unk>',
						    'bos_token': '<s>',
						    'eos_token': '</s>',
						}
						A_						:							int      = {
						    'feature_size': 1,
						    'padding_value': 0.0,
						    'sampling_rate': 1_6_0_0_0,
						    'return_attention_mask': False,
						    'do_normalize': True,
						}

						A_						:							Tuple      = tempfile.mkdtemp()
						A_						:							Any      = os.path.join(self.tmpdirname						,      VOCAB_FILES_NAMES['vocab_file']   )
						A_						:							int      = os.path.join(self.tmpdirname						,      lowercase   )
						with open(self.vocab_file						,      'w'						,      encoding='utf-8'   ) as fp:
								fp.write(json.dumps(lowercase   ) + '\n'   )

						with open(self.feature_extraction_file						,      'w'						,      encoding='utf-8'   ) as fp:
								fp.write(json.dumps(lowercase   ) + '\n'   )

						# load decoder from hub
						A_						:							Dict      = 'hf-internal-testing/ngram-beam-search-decoder'
				def     lowerCAmelCase_  (		self						,      **lowercase   ):



						"""simple docstring"""





						A_						:							List[str]      = self.add_kwargs_tokens_map.copy()
						kwargs.update(lowercase   )
						return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname						,      **lowercase   )
				def     lowerCAmelCase_  (		self						,      **lowercase   ):



						"""simple docstring"""





						return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname						,      **lowercase   )
				def     lowerCAmelCase_  (		self						,      **lowercase   ):



						"""simple docstring"""





						return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name						,      **lowercase   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						shutil.rmtree(self.tmpdirname   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Optional[Any]      = self.get_tokenizer()
						A_						:							List[str]      = self.get_feature_extractor()
						A_						:							int      = self.get_decoder()

						A_						:							Union[str, Any]      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						processor.save_pretrained(self.tmpdirname   )
						A_						:							Union[str, Any]      = WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname   )

						# tokenizer
						self.assertEqual(processor.tokenizer.get_vocab()						,      tokenizer.get_vocab()   )
						self.assertIsInstance(processor.tokenizer						,      lowercase   )

						# feature extractor
						self.assertEqual(processor.feature_extractor.to_json_string()						,      feature_extractor.to_json_string()   )
						self.assertIsInstance(processor.feature_extractor						,      lowercase   )

						# decoder
						self.assertEqual(processor.decoder._alphabet.labels						,      decoder._alphabet.labels   )
						self.assertEqual(
						    processor.decoder.model_container[decoder._model_key]._unigram_set						,      decoder.model_container[decoder._model_key]._unigram_set						,      )
						self.assertIsInstance(processor.decoder						,      lowercase   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							str      = WavaVecaProcessorWithLM(
						    tokenizer=self.get_tokenizer()						,      feature_extractor=self.get_feature_extractor()						,      decoder=self.get_decoder()   )
						processor.save_pretrained(self.tmpdirname   )

						# make sure that error is thrown when decoder alphabet doesn't match
						A_						:							Optional[int]      = WavaVecaProcessorWithLM.from_pretrained(
						    self.tmpdirname						,      alpha=5.0						,      beta=3.0						,      score_boundary=-7.0						,      unk_score_offset=3   )

						# decoder
						self.assertEqual(processor.language_model.alpha						,      5.0   )
						self.assertEqual(processor.language_model.beta						,      3.0   )
						self.assertEqual(processor.language_model.score_boundary						,      -7.0   )
						self.assertEqual(processor.language_model.unk_score_offset						,      3   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Union[str, Any]      = self.get_tokenizer()
						# add token to trigger raise
						tokenizer.add_tokens(['xx']   )
						with self.assertRaisesRegex(lowercase						,      'include'   ):
								WavaVecaProcessorWithLM(
								    tokenizer=lowercase						,      feature_extractor=self.get_feature_extractor()						,      decoder=self.get_decoder()   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Dict      = self.get_feature_extractor()
						A_						:							Any      = self.get_tokenizer()
						A_						:							int      = self.get_decoder()

						A_						:							List[Any]      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						A_						:							Tuple      = floats_list((3, 1_0_0_0)   )

						A_						:							str      = feature_extractor(lowercase						,      return_tensors='np'   )
						A_						:							Union[str, Any]      = processor(lowercase						,      return_tensors='np'   )

						for key in input_feat_extract.keys():
								self.assertAlmostEqual(input_feat_extract[key].sum()						,      input_processor[key].sum()						,      delta=1E-2   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Tuple      = self.get_feature_extractor()
						A_						:							Any      = self.get_tokenizer()
						A_						:							Optional[Any]      = self.get_decoder()

						A_						:							Optional[Any]      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						A_						:							List[Any]      = 'This is a test string'

						A_						:							Optional[Any]      = processor(text=lowercase   )

						A_						:							List[Any]      = tokenizer(lowercase   )

						for key in encoded_tok.keys():
								self.assertListEqual(encoded_tok[key]						,      encoded_processor[key]   )
				def     lowerCAmelCase_  (		self						,      lowercase=(2, 1_0, 1_6)						,      lowercase=7_7   ):



						"""simple docstring"""





						np.random.seed(lowercase   )
						return np.random.rand(*lowercase   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Optional[int]      = self.get_feature_extractor()
						A_						:							str      = self.get_tokenizer()
						A_						:							Any      = self.get_decoder()

						A_						:							Tuple      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						A_						:							str      = self._get_dummy_logits(shape=(1_0, 1_6)						,      seed=1_3   )

						A_						:							Optional[int]      = processor.decode(lowercase   )

						A_						:							Optional[Any]      = decoder.decode_beams(lowercase   )[0]

						self.assertEqual(decoded_decoder[0]						,      decoded_processor.text   )
						self.assertEqual('</s> <s> </s>'						,      decoded_processor.text   )
						self.assertEqual(decoded_decoder[-2]						,      decoded_processor.logit_score   )
						self.assertEqual(decoded_decoder[-1]						,      decoded_processor.lm_score   )
				@parameterized.expand([[None], ['fork'], ['spawn']]   )
				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						A_						:							str      = self.get_feature_extractor()
						A_						:							Optional[Any]      = self.get_tokenizer()
						A_						:							List[str]      = self.get_decoder()

						A_						:							List[Any]      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						A_						:							List[Any]      = self._get_dummy_logits()

						# note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
						#       otherwise, the LM won't be available to the pool's sub-processes.
						# manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
						if pool_context is None:
								A_						:							Tuple      = processor.batch_decode(lowercase   )
						else:
								with get_context(lowercase   ).Pool() as pool:
										A_						:							int      = processor.batch_decode(lowercase						,      lowercase   )

						A_						:							int      = list(lowercase   )

						with get_context('fork'   ).Pool() as p:
								A_						:							List[Any]      = decoder.decode_beams_batch(lowercase						,      lowercase   )

						A_  ,     A_  ,     A_						:							int      = [], [], []
						for beams in decoded_beams:
								texts_decoder.append(beams[0][0]   )
								logit_scores_decoder.append(beams[0][-2]   )
								lm_scores_decoder.append(beams[0][-1]   )

						self.assertListEqual(lowercase						,      decoded_processor.text   )
						self.assertListEqual(['<s> <s> </s>', '<s> <s> <s>']						,      decoded_processor.text   )
						self.assertListEqual(lowercase						,      decoded_processor.logit_score   )
						self.assertListEqual(lowercase						,      decoded_processor.lm_score   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							List[Any]      = self.get_feature_extractor()
						A_						:							str      = self.get_tokenizer()
						A_						:							Optional[Any]      = self.get_decoder()

						A_						:							str      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						A_						:							Optional[int]      = self._get_dummy_logits()

						A_						:							Any      = 1_5
						A_						:							Optional[int]      = -20.0
						A_						:							Tuple      = -4.0

						A_						:							str      = processor.batch_decode(
						    lowercase						,      beam_width=lowercase						,      beam_prune_logp=lowercase						,      token_min_logp=lowercase						,      )
						A_						:							Union[str, Any]      = decoded_processor_out.text

						A_						:							List[Any]      = list(lowercase   )

						with get_context('fork'   ).Pool() as pool:
								A_						:							List[Any]      = decoder.decode_beams_batch(
								    lowercase						,      lowercase						,      beam_width=lowercase						,      beam_prune_logp=lowercase						,      token_min_logp=lowercase						,      )

						A_						:							Tuple      = [d[0][0] for d in decoded_decoder_out]
						A_						:							int      = [d[0][2] for d in decoded_decoder_out]
						A_						:							int      = [d[0][3] for d in decoded_decoder_out]

						self.assertListEqual(lowercase						,      lowercase   )
						self.assertListEqual(['</s> <s> <s>', '<s> <s> <s>']						,      lowercase   )

						self.assertTrue(np.array_equal(lowercase						,      decoded_processor_out.logit_score   )   )
						self.assertTrue(np.allclose([-20.054, -18.447]						,      lowercase						,      atol=1E-3   )   )

						self.assertTrue(np.array_equal(lowercase						,      decoded_processor_out.lm_score   )   )
						self.assertTrue(np.allclose([-15.554, -13.9474]						,      lowercase						,      atol=1E-3   )   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Tuple      = self.get_feature_extractor()
						A_						:							Any      = self.get_tokenizer()
						A_						:							Optional[Any]      = self.get_decoder()

						A_						:							Optional[Any]      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						A_						:							Dict      = self._get_dummy_logits()

						A_						:							str      = 2.0
						A_						:							List[Any]      = 5.0
						A_						:							Optional[int]      = -20.0
						A_						:							List[Any]      = True

						A_						:							List[str]      = processor.batch_decode(
						    lowercase						,      alpha=lowercase						,      beta=lowercase						,      unk_score_offset=lowercase						,      lm_score_boundary=lowercase						,      )
						A_						:							List[Any]      = decoded_processor_out.text

						A_						:							List[str]      = list(lowercase   )
						decoder.reset_params(
						    alpha=lowercase						,      beta=lowercase						,      unk_score_offset=lowercase						,      lm_score_boundary=lowercase						,      )

						with get_context('fork'   ).Pool() as pool:
								A_						:							Union[str, Any]      = decoder.decode_beams_batch(
								    lowercase						,      lowercase						,      )

						A_						:							Any      = [d[0][0] for d in decoded_decoder_out]

						self.assertListEqual(lowercase						,      lowercase   )
						self.assertListEqual(['<s> </s> <s> </s> </s>', '</s> </s> <s> </s> </s>']						,      lowercase   )
						A_						:							List[str]      = processor.decoder.model_container[processor.decoder._model_key]
						self.assertEqual(lm_model.alpha						,      2.0   )
						self.assertEqual(lm_model.beta						,      5.0   )
						self.assertEqual(lm_model.unk_score_offset						,      -20.0   )
						self.assertEqual(lm_model.score_boundary						,      lowercase   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Optional[Any]      = WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'   )

						A_						:							Optional[Any]      = processor.decoder.model_container[processor.decoder._model_key]
						A_						:							Dict      = Path(language_model._kenlm_model.path.decode('utf-8'   )   ).parent.parent.absolute()

						A_						:							int      = os.listdir(lowercase   )
						A_						:							Tuple      = ['alphabet.json', 'language_model']

						downloaded_decoder_files.sort()
						expected_decoder_files.sort()

						# test that only decoder relevant files from
						# https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
						# are downloaded and none of the rest (e.g. README.md, ...)
						self.assertListEqual(lowercase						,      lowercase   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Optional[int]      = snapshot_download('hf-internal-testing/processor_with_lm'   )

						A_						:							str      = WavaVecaProcessorWithLM.from_pretrained(lowercase   )

						A_						:							Optional[int]      = processor.decoder.model_container[processor.decoder._model_key]
						A_						:							int      = Path(language_model._kenlm_model.path.decode('utf-8'   )   ).parent.parent.absolute()

						A_						:							int      = os.listdir(lowercase   )
						A_						:							Any      = os.listdir(lowercase   )

						local_decoder_files.sort()
						expected_decoder_files.sort()

						# test that both decoder form hub and local files in cache are the same
						self.assertListEqual(lowercase						,      lowercase   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							int      = WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'   )
						A_						:							Dict      = AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'   )

						A_						:							List[Any]      = floats_list((3, 1_0_0_0)   )

						A_						:							Any      = processor_wavaveca(lowercase						,      return_tensors='np'   )
						A_						:							List[str]      = processor_auto(lowercase						,      return_tensors='np'   )

						for key in input_wavaveca.keys():
								self.assertAlmostEqual(input_wavaveca[key].sum()						,      input_auto[key].sum()						,      delta=1E-2   )

						A_						:							Any      = self._get_dummy_logits()

						A_						:							Tuple      = processor_wavaveca.batch_decode(lowercase   )
						A_						:							Optional[int]      = processor_auto.batch_decode(lowercase   )

						self.assertListEqual(decoded_wavaveca.text						,      decoded_auto.text   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Any      = self.get_feature_extractor()
						A_						:							Any      = self.get_tokenizer()
						A_						:							Optional[int]      = self.get_decoder()

						A_						:							List[Any]      = WavaVecaProcessorWithLM(tokenizer=lowercase						,      feature_extractor=lowercase						,      decoder=lowercase   )

						self.assertListEqual(
						    processor.model_input_names						,      feature_extractor.model_input_names						,      msg='`processor` and `feature_extractor` model input names do not match'						,      )
				@staticmethod
				def     lowerCAmelCase_  (		lowercase						,      lowercase   ):



						"""simple docstring"""





						A_						:							Dict      = [d[key] for d in offsets]
						return retrieved_list
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Any      = WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'   )
						A_						:							Optional[Any]      = self._get_dummy_logits()[0]

						A_						:							Optional[Any]      = processor.decode(lowercase						,      output_word_offsets=lowercase   )
						# check Wav2Vec2CTCTokenizerOutput keys for word
						self.assertEqual(len(outputs.keys()   )						,      4   )
						self.assertTrue('text' in outputs   )
						self.assertTrue('word_offsets' in outputs   )
						self.assertTrue(isinstance(lowercase						,      lowercase   )   )

						self.assertEqual(' '.join(self.get_from_offsets(outputs['word_offsets']						,      'word'   )   )						,      outputs.text   )
						self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,      'word'   )						,      ['<s>', '<s>', '</s>']   )
						self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,      'start_offset'   )						,      [0, 2, 4]   )
						self.assertListEqual(self.get_from_offsets(outputs['word_offsets']						,      'end_offset'   )						,      [1, 3, 5]   )
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Union[str, Any]      = WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'   )
						A_						:							Optional[Any]      = self._get_dummy_logits()

						A_						:							List[Any]      = processor.batch_decode(lowercase						,      output_word_offsets=lowercase   )

						# check Wav2Vec2CTCTokenizerOutput keys for word
						self.assertEqual(len(outputs.keys()   )						,      4   )
						self.assertTrue('text' in outputs   )
						self.assertTrue('word_offsets' in outputs   )
						self.assertTrue(isinstance(lowercase						,      lowercase   )   )

						self.assertListEqual(
						    [' '.join(self.get_from_offsets(lowercase						,      'word'   )   ) for o in outputs['word_offsets']]						,      outputs.text   )
						self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,      'word'   )						,      ['<s>', '<s>', '</s>']   )
						self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,      'start_offset'   )						,      [0, 2, 4]   )
						self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]						,      'end_offset'   )						,      [1, 3, 5]   )







				@slow
				@require_torch
				@require_torchaudio
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						import torch

						A_						:							Any      = load_dataset('common_voice'						,      'en'						,      split='train'						,      streaming=lowercase   )
						A_						:							List[str]      = ds.cast_column('audio'						,      datasets.Audio(sampling_rate=1_6_0_0_0   )   )
						A_						:							Dict      = iter(lowercase   )
						A_						:							List[str]      = next(lowercase   )

						A_						:							int      = AutoProcessor.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'   )
						A_						:							List[str]      = WavaVecaForCTC.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'   )

						# compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
						A_						:							Optional[Any]      = processor(sample['audio']['array']						,      return_tensors='pt'   ).input_values

						with torch.no_grad():
								A_						:							Optional[int]      = model(lowercase   ).logits.cpu().numpy()

						A_						:							Dict      = processor.decode(logits[0]						,      output_word_offsets=lowercase   )

						A_						:							Optional[Any]      = model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
						A_						:							List[Any]      = [
						    {
						        'start_time': d['start_offset'] * time_offset,
						        'end_time': d['end_offset'] * time_offset,
						        'word': d['word'],
						    }
						    for d in output['word_offsets']
						]

						A_						:							Tuple      = 'WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL'

						# output words
						self.assertEqual(' '.join(self.get_from_offsets(lowercase						,      'word'   )   )						,      lowercase   )
						self.assertEqual(' '.join(self.get_from_offsets(lowercase						,      'word'   )   )						,      output.text   )

						# output times
						A_						:							Dict      = torch.tensor(self.get_from_offsets(lowercase						,      'start_time'   )   )
						A_						:							Optional[Any]      = torch.tensor(self.get_from_offsets(lowercase						,      'end_time'   )   )

						# fmt: off
						A_						:							Optional[Any]      = torch.tensor([1.4199, 1.6599, 2.2599, 3.0, 3.24, 3.5999, 3.7999, 4.0999, 4.26, 4.94, 5.28, 5.6599, 5.78, 5.94, 6.32, 6.5399, 6.6599]   )
						A_						:							Optional[Any]      = torch.tensor([1.5399, 1.8999, 2.9, 3.16, 3.5399, 3.72, 4.0199, 4.1799, 4.76, 5.1599, 5.5599, 5.6999, 5.86, 6.1999, 6.38, 6.6199, 6.94]   )
						# fmt: on

						self.assertTrue(torch.allclose(lowercase						,      lowercase						,      atol=0.01   )   )
						self.assertTrue(torch.allclose(lowercase						,      lowercase						,      atol=0.01   )   )







from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


_UpperCAmelCase          =						logging.get_logger(__name__)

# General docstring
_UpperCAmelCase          =						"""RegNetConfig"""

# Base docstring
_UpperCAmelCase          =						"""facebook/regnet-y-040"""
_UpperCAmelCase          =						[1, 1088, 7, 7]

# Image classification docstring
_UpperCAmelCase          =						"""facebook/regnet-y-040"""
_UpperCAmelCase          =						"""tabby, tabby cat"""

_UpperCAmelCase          =						[
    """facebook/regnet-y-040""",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase = 3						,      lowercase = 1						,      lowercase = 1						,      lowercase = "relu"						,      **lowercase						,      ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						# The padding and conv has been verified in
						# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
						A_						:							List[str]      = tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2   )
						A_						:							Tuple      = tf.keras.layers.ConvaD(
						    filters=lowercase						,      kernel_size=lowercase						,      strides=lowercase						,      padding='VALID'						,      groups=lowercase						,      use_bias=lowercase						,      name='convolution'						,      )
						A_						:							Dict      = tf.keras.layers.BatchNormalization(epsilon=1E-5						,      momentum=0.9						,      name='normalization'   )
						A_						:							Optional[int]      = ACTaFN[activation] if activation is not None else tf.identity







				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						A_						:							int      = self.convolution(self.padding(lowercase   )   )
						A_						:							Optional[Any]      = self.normalization(lowercase   )
						A_						:							Union[str, Any]      = self.activation(lowercase   )
						return hidden_state





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							Dict      = config.num_channels
						A_						:							Dict      = TFRegNetConvLayer(
						    out_channels=config.embedding_size						,      kernel_size=3						,      stride=2						,      activation=config.hidden_act						,      name='embedder'						,      )







				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						A_						:							Union[str, Any]      = shape_list(lowercase   )[1]
						if tf.executing_eagerly() and num_channels != self.num_channels:
								raise ValueError(
								    'Make sure that the channel dimension of the pixel values match with the one set in the configuration.'   )

						# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
						# So change the input format from `NCHW` to `NHWC`.
						# shape = (batch_size, in_height, in_width, in_channels=num_channels)
						A_						:							Optional[Any]      = tf.transpose(lowercase						,      perm=(0, 2, 3, 1)   )
						A_						:							Dict      = self.embedder(lowercase   )
						return hidden_state





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase = 2						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							Optional[int]      = tf.keras.layers.ConvaD(
						    filters=lowercase						,      kernel_size=1						,      strides=lowercase						,      use_bias=lowercase						,      name='convolution'   )
						A_						:							Tuple      = tf.keras.layers.BatchNormalization(epsilon=1E-5						,      momentum=0.9						,      name='normalization'   )







				def     lowerCAmelCase_  (		self						,      lowercase						,      lowercase = False   ):



						"""simple docstring"""





						return self.normalization(self.convolution(lowercase   )						,      training=lowercase   )





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							List[Any]      = tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase						,      name='pooler'   )
						A_						:							Any      = [
						    tf.keras.layers.ConvaD(filters=lowercase						,      kernel_size=1						,      activation='relu'						,      name='attention.0'   ),
						    tf.keras.layers.ConvaD(filters=lowercase						,      kernel_size=1						,      activation='sigmoid'						,      name='attention.2'   ),
						]







				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						A_						:							str      = self.pooler(lowercase   )
						for layer_module in self.attention:
								A_						:							List[str]      = layer_module(lowercase   )
						A_						:							str      = hidden_state * pooled
						return hidden_state





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase						,      lowercase						,      lowercase = 1						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							Union[str, Any]      = in_channels != out_channels or stride != 1
						A_						:							Optional[Any]      = max(1						,      out_channels // config.groups_width   )
						A_						:							List[Any]      = (
						    TFRegNetShortCut(lowercase						,      stride=lowercase						,      name='shortcut'   )
						    if should_apply_shortcut
						    else tf.keras.layers.Activation('linear'						,      name='shortcut'   )
						)
						# `self.layers` instead of `self.layer` because that is a reserved argument.
						A_						:							Optional[Any]      = [
						    TFRegNetConvLayer(lowercase						,      kernel_size=1						,      activation=config.hidden_act						,      name='layer.0'   ),
						    TFRegNetConvLayer(
						        lowercase						,      stride=lowercase						,      groups=lowercase						,      activation=config.hidden_act						,      name='layer.1'   ),
						    TFRegNetConvLayer(lowercase						,      kernel_size=1						,      activation=lowercase						,      name='layer.2'   ),
						]
						A_						:							int      = ACTaFN[config.hidden_act]







				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						A_						:							Optional[Any]      = hidden_state
						for layer_module in self.layers:
								A_						:							Union[str, Any]      = layer_module(lowercase   )
						A_						:							int      = self.shortcut(lowercase   )
						hidden_state += residual
						A_						:							Optional[int]      = self.activation(lowercase   )
						return hidden_state





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase						,      lowercase						,      lowercase = 1						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							Any      = in_channels != out_channels or stride != 1
						A_						:							Union[str, Any]      = max(1						,      out_channels // config.groups_width   )
						A_						:							str      = (
						    TFRegNetShortCut(lowercase						,      stride=lowercase						,      name='shortcut'   )
						    if should_apply_shortcut
						    else tf.keras.layers.Activation('linear'						,      name='shortcut'   )
						)
						A_						:							Optional[int]      = [
						    TFRegNetConvLayer(lowercase						,      kernel_size=1						,      activation=config.hidden_act						,      name='layer.0'   ),
						    TFRegNetConvLayer(
						        lowercase						,      stride=lowercase						,      groups=lowercase						,      activation=config.hidden_act						,      name='layer.1'   ),
						    TFRegNetSELayer(lowercase						,      reduced_channels=int(round(in_channels / 4   )   )						,      name='layer.2'   ),
						    TFRegNetConvLayer(lowercase						,      kernel_size=1						,      activation=lowercase						,      name='layer.3'   ),
						]
						A_						:							str      = ACTaFN[config.hidden_act]







				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						A_						:							Any      = hidden_state
						for layer_module in self.layers:
								A_						:							Optional[Any]      = layer_module(lowercase   )
						A_						:							Optional[Any]      = self.shortcut(lowercase   )
						hidden_state += residual
						A_						:							Tuple      = self.activation(lowercase   )
						return hidden_state





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase						,      lowercase						,      lowercase = 2						,      lowercase = 2						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )

						A_						:							List[Any]      = TFRegNetXLayer if config.layer_type == 'x' else TFRegNetYLayer
						A_						:							int      = [
						    # downsampling is done in the first layer with stride of 2
						    layer(lowercase						,      lowercase						,      lowercase						,      stride=lowercase						,      name='layers.0'   ),
						    *[layer(lowercase						,      lowercase						,      lowercase						,      name=F'''layers.{i+1}'''   ) for i in range(depth - 1   )],
						]







				def     lowerCAmelCase_  (		self						,      lowercase   ):



						"""simple docstring"""





						for layer_module in self.layers:
								A_						:							List[str]      = layer_module(lowercase   )
						return hidden_state





class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							Any      = []
						# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
						self.stages.append(
						    TFRegNetStage(
						        lowercase						,      config.embedding_size						,      config.hidden_sizes[0]						,      stride=2 if config.downsample_in_first_stage else 1						,      depth=config.depths[0]						,      name='stages.0'						,      )   )
						A_						:							int      = zip(config.hidden_sizes						,      config.hidden_sizes[1:]   )
						for i, ((in_channels, out_channels), depth) in enumerate(zip(lowercase						,      config.depths[1:]   )   ):
								self.stages.append(TFRegNetStage(lowercase						,      lowercase						,      lowercase						,      depth=lowercase						,      name=F'''stages.{i+1}'''   )   )







				def     lowerCAmelCase_  (		self						,      lowercase						,      lowercase = False						,      lowercase = True   ):



						"""simple docstring"""





						A_						:							Optional[Any]      = () if output_hidden_states else None

						for stage_module in self.stages:
								if output_hidden_states:
										A_						:							Optional[Any]      = hidden_states + (hidden_state,)

								A_						:							List[Any]      = stage_module(lowercase   )

						if output_hidden_states:
								A_						:							str      = hidden_states + (hidden_state,)

						if not return_dict:
								return tuple(v for v in [hidden_state, hidden_states] if v is not None   )

						return TFBaseModelOutputWithNoAttention(last_hidden_state=lowercase						,      hidden_states=lowercase   )





@keras_serializable
class UpperCAmelCase		(   tf.keras.layers.Layer    ):


				'''simple docstring'''






				lowerCamelCase_        =   RegNetConfig
				def __init__(		self						,      lowercase						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(**lowercase   )
						A_						:							List[Any]      = config
						A_						:							List[str]      = TFRegNetEmbeddings(lowercase						,      name='embedder'   )
						A_						:							Dict      = TFRegNetEncoder(lowercase						,      name='encoder'   )
						A_						:							Any      = tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase						,      name='pooler'   )







				@unpack_inputs
				def     lowerCAmelCase_  (		self						,      lowercase						,      lowercase = None						,      lowercase = None						,      lowercase = False						,      ):



						"""simple docstring"""





						A_						:							Dict      = (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						A_						:							Any      = return_dict if return_dict is not None else self.config.use_return_dict

						A_						:							int      = self.embedder(lowercase						,      training=lowercase   )

						A_						:							int      = self.encoder(
						    lowercase						,      output_hidden_states=lowercase						,      return_dict=lowercase						,      training=lowercase   )

						A_						:							List[Any]      = encoder_outputs[0]
						A_						:							Any      = self.pooler(lowercase   )

						# Change to NCHW output format have uniformity in the modules
						A_						:							Optional[int]      = tf.transpose(lowercase						,      perm=(0, 3, 1, 2)   )
						A_						:							List[str]      = tf.transpose(lowercase						,      perm=(0, 3, 1, 2)   )

						# Change the other hidden state outputs to NCHW as well
						if output_hidden_states:
								A_						:							Dict      = tuple([tf.transpose(lowercase						,      perm=(0, 3, 1, 2)   ) for h in encoder_outputs[1]]   )

						if not return_dict:
								return (last_hidden_state, pooled_output) + encoder_outputs[1:]

						return TFBaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=lowercase						,      pooler_output=lowercase						,      hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states						,      )





class UpperCAmelCase		(   __A    ):


				'''simple docstring'''






				lowerCamelCase_        =   RegNetConfig
				lowerCamelCase_        =   '''regnet'''
				lowerCamelCase_        =   '''pixel_values'''
				@property
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4)						,      dtype=tf.floataa   )}

_UpperCAmelCase          =						r"""
    Parameters:
    This model is a Tensorflow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and
    behavior.
        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""

_UpperCAmelCase          =						r"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConveNextImageProcessor.__call__`] for details.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""





@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.'''   ,			__A   ,			)
class UpperCAmelCase		(   __A    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      *lowercase						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(lowercase						,      *lowercase						,      **lowercase   )
						A_						:							Union[str, Any]      = TFRegNetMainLayer(lowercase						,      name='regnet'   )







				@unpack_inputs
				@add_start_docstrings_to_model_forward(lowercase   )
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC						,      output_type=lowercase						,      config_class=_CONFIG_FOR_DOC						,      modality='vision'						,      expected_output=_EXPECTED_OUTPUT_SHAPE						,      )
				def     lowerCAmelCase_  (		self						,      lowercase						,      lowercase = None						,      lowercase = None						,      lowercase=False						,      ):



						"""simple docstring"""





						A_						:							Union[str, Any]      = (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						A_						:							str      = return_dict if return_dict is not None else self.config.use_return_dict

						A_						:							int      = self.regnet(
						    pixel_values=lowercase						,      output_hidden_states=lowercase						,      return_dict=lowercase						,      training=lowercase						,      )
						if not return_dict:
								return (outputs[0],) + outputs[1:]

						return TFBaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=outputs.last_hidden_state						,      pooler_output=outputs.pooler_output						,      hidden_states=outputs.hidden_states						,      )





@add_start_docstrings(
    '''
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''   ,			__A   ,			)
class UpperCAmelCase		(   __A   ,			__A    ):


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      *lowercase						,      **lowercase   ):



						"""simple docstring"""





						super().__init__(lowercase						,      *lowercase						,      **lowercase   )
						A_						:							str      = config.num_labels
						A_						:							Optional[Any]      = TFRegNetMainLayer(lowercase						,      name='regnet'   )
						# classification head
						A_						:							int      = [
						    tf.keras.layers.Flatten(),
						    tf.keras.layers.Dense(config.num_labels						,      name='classifier.1'   ) if config.num_labels > 0 else tf.identity,
						]







				@unpack_inputs
				@add_start_docstrings_to_model_forward(lowercase   )
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT						,      output_type=lowercase						,      config_class=_CONFIG_FOR_DOC						,      expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT						,      )
				def     lowerCAmelCase_  (		self						,      lowercase = None						,      lowercase = None						,      lowercase = None						,      lowercase = None						,      lowercase=False						,      ):



						"""simple docstring"""





						A_						:							Any      = (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						A_						:							Union[str, Any]      = return_dict if return_dict is not None else self.config.use_return_dict

						A_						:							Dict      = self.regnet(
						    lowercase						,      output_hidden_states=lowercase						,      return_dict=lowercase						,      training=lowercase   )

						A_						:							Any      = outputs.pooler_output if return_dict else outputs[1]

						A_						:							Union[str, Any]      = self.classifier[0](lowercase   )
						A_						:							Dict      = self.classifier[1](lowercase   )

						A_						:							Dict      = None if labels is None else self.hf_compute_loss(labels=lowercase						,      logits=lowercase   )

						if not return_dict:
								A_						:							List[str]      = (logits,) + outputs[2:]
								return ((loss,) + output) if loss is not None else output

						return TFSequenceClassifierOutput(loss=lowercase						,      logits=lowercase						,      hidden_states=outputs.hidden_states   )






