


"""simple docstring"""

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  lowerCamelCase__    (		unittest.TestCase     ):
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]				):






       '''simple docstring'''


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

       __lowercase						=						{
           """unk_token""": """<unk>""",
           """bos_token""": """<s>""",
           """eos_token""": """</s>""",
       }
       __lowercase						=						{
           """feature_size""": 1,
           """padding_value""": 0.0,
           """sampling_rate""": 1_6_0_0_0,
           """return_attention_mask""": False,
           """do_normalize""": True,
       }

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

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

       # load decoder from hub
       __lowercase						=						"""hf-internal-testing/ngram-beam-search-decoder"""
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							**A_						:						List[str]				):






       '''simple docstring'''


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






       '''simple docstring'''


       return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname   ,							**A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							**A_						:						Union[str, Any]				):






       '''simple docstring'''


       return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name   ,							**A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       shutil.rmtree(self.tmpdirname				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       processor.save_pretrained(self.tmpdirname				)
       __lowercase						=						WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname				)

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

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

       # 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   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						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
       __lowercase						=						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 SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						self.get_tokenizer()
       # add token to trigger raise
       tokenizer.add_tokens(["""xx"""]				)
       with self.assertRaisesRegex(A_   ,							"""include"""				):
        WavaVecaProcessorWithLM(
            tokenizer=A_   ,							feature_extractor=self.get_feature_extractor()   ,							decoder=self.get_decoder()				)
      def SCREAMING_SNAKE_CASE_						(		self						:						int				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       __lowercase						=						floats_list((3, 1_0_0_0)				)

       __lowercase						=						feature_extractor(A_   ,							return_tensors="""np"""				)
       __lowercase						=						processor(A_   ,							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 SCREAMING_SNAKE_CASE_						(		self						:						int				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       __lowercase						=						"""This is a test string"""

       __lowercase						=						processor(text=A_				)

       __lowercase						=						tokenizer(A_				)

       for key in encoded_tok.keys():
        self.assertListEqual(encoded_tok[key]   ,							encoded_processor[key]				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						List[Any]=(2, 1_0, 1_6)   ,							A_						:						Union[str, Any]=7_7				):






       '''simple docstring'''


       np.random.seed(A_				)
       return np.random.rand(*A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       __lowercase						=						self._get_dummy_logits(shape=(1_0, 1_6)   ,							seed=1_3				)

       __lowercase						=						processor.decode(A_				)

       __lowercase						=						decoder.decode_beams(A_				)[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 SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]   ,							A_						:						Optional[Any]				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       __lowercase						=						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:
        __lowercase						=						processor.batch_decode(A_				)
       else:
        with get_context(A_				).Pool() as pool:
         __lowercase						=						processor.batch_decode(A_   ,							A_				)

       __lowercase						=						list(A_				)

       with get_context("""fork"""				).Pool() as p:
        __lowercase						=						decoder.decode_beams_batch(A_   ,							A_				)

       __lowercase				,		__lowercase				,		__lowercase						=						[], [], []
       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(A_   ,							decoded_processor.text				)
       self.assertListEqual(["""<s> <s> </s>""", """<s> <s> <s>"""]   ,							decoded_processor.text				)
       self.assertListEqual(A_   ,							decoded_processor.logit_score				)
       self.assertListEqual(A_   ,							decoded_processor.lm_score				)
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       __lowercase						=						self._get_dummy_logits()

       __lowercase						=						1_5
       __lowercase						=						-20.0
       __lowercase						=						-4.0

       __lowercase						=						processor.batch_decode(
           A_   ,							beam_width=A_   ,							beam_prune_logp=A_   ,							token_min_logp=A_   ,							)
       __lowercase						=						decoded_processor_out.text

       __lowercase						=						list(A_				)

       with get_context("""fork"""				).Pool() as pool:
        __lowercase						=						decoder.decode_beams_batch(
            A_   ,							A_   ,							beam_width=A_   ,							beam_prune_logp=A_   ,							token_min_logp=A_   ,							)

       __lowercase						=						[d[0][0] for d in decoded_decoder_out]
       __lowercase						=						[d[0][2] for d in decoded_decoder_out]
       __lowercase						=						[d[0][3] for d in decoded_decoder_out]

       self.assertListEqual(A_   ,							A_				)
       self.assertListEqual(["""</s> <s> <s>""", """<s> <s> <s>"""]   ,							A_				)

       self.assertTrue(np.array_equal(A_   ,							decoded_processor_out.logit_score				)				)
       self.assertTrue(np.allclose([-20.0_54, -18.4_47]   ,							A_   ,							atol=1e-3				)				)

       self.assertTrue(np.array_equal(A_   ,							decoded_processor_out.lm_score				)				)
       self.assertTrue(np.allclose([-15.5_54, -13.94_74]   ,							A_   ,							atol=1e-3				)				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       __lowercase						=						self._get_dummy_logits()

       __lowercase						=						2.0
       __lowercase						=						5.0
       __lowercase						=						-20.0
       __lowercase						=						True

       __lowercase						=						processor.batch_decode(
           A_   ,							alpha=A_   ,							beta=A_   ,							unk_score_offset=A_   ,							lm_score_boundary=A_   ,							)
       __lowercase						=						decoded_processor_out.text

       __lowercase						=						list(A_				)
       decoder.reset_params(
           alpha=A_   ,							beta=A_   ,							unk_score_offset=A_   ,							lm_score_boundary=A_   ,							)

       with get_context("""fork"""				).Pool() as pool:
        __lowercase						=						decoder.decode_beams_batch(
            A_   ,							A_   ,							)

       __lowercase						=						[d[0][0] for d in decoded_decoder_out]

       self.assertListEqual(A_   ,							A_				)
       self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""]   ,							A_				)
       __lowercase						=						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   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						int				):






       '''simple docstring'''


       __lowercase						=						WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)

       __lowercase						=						processor.decoder.model_container[processor.decoder._model_key]
       __lowercase						=						Path(language_model._kenlm_model.path.decode("""utf-8"""				)				).parent.parent.absolute()

       __lowercase						=						os.listdir(A_				)
       __lowercase						=						["""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(A_   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						snapshot_download("""hf-internal-testing/processor_with_lm"""				)

       __lowercase						=						WavaVecaProcessorWithLM.from_pretrained(A_				)

       __lowercase						=						processor.decoder.model_container[processor.decoder._model_key]
       __lowercase						=						Path(language_model._kenlm_model.path.decode("""utf-8"""				)				).parent.parent.absolute()

       __lowercase						=						os.listdir(A_				)
       __lowercase						=						os.listdir(A_				)

       local_decoder_files.sort()
       expected_decoder_files.sort()

       # test that both decoder form hub and local files in cache are the same
       self.assertListEqual(A_   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]				):






       '''simple docstring'''


       __lowercase						=						WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)
       __lowercase						=						AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""				)

       __lowercase						=						floats_list((3, 1_0_0_0)				)

       __lowercase						=						processor_wavaveca(A_   ,							return_tensors="""np"""				)
       __lowercase						=						processor_auto(A_   ,							return_tensors="""np"""				)

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

       __lowercase						=						self._get_dummy_logits()

       __lowercase						=						processor_wavaveca.batch_decode(A_				)
       __lowercase						=						processor_auto.batch_decode(A_				)

       self.assertListEqual(decoded_wavaveca.text   ,							decoded_auto.text				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						self.get_feature_extractor()
       __lowercase						=						self.get_tokenizer()
       __lowercase						=						self.get_decoder()

       __lowercase						=						WavaVecaProcessorWithLM(tokenizer=A_   ,							feature_extractor=A_   ,							decoder=A_				)

       self.assertListEqual(
           processor.model_input_names   ,							feature_extractor.model_input_names   ,							msg="""`processor` and `feature_extractor` model input names do not match"""   ,							)
      @staticmethod
      def SCREAMING_SNAKE_CASE_						(		A_						:						List[Any]   ,							A_						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						[d[key] for d in offsets]
       return retrieved_list
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)
       __lowercase						=						self._get_dummy_logits()[0]

       __lowercase						=						processor.decode(A_   ,							output_word_offsets=A_				)
       # 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(A_   ,							A_				)				)

       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 SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)
       __lowercase						=						self._get_dummy_logits()

       __lowercase						=						processor.batch_decode(A_   ,							output_word_offsets=A_				)

       # 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(A_   ,							A_				)				)

       self.assertListEqual(
           [""" """.join(self.get_from_offsets(A_   ,							"""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 SCREAMING_SNAKE_CASE_						(		self						:						Any				):






       '''simple docstring'''


       import torch

       __lowercase						=						load_dataset("""common_voice"""   ,							"""en"""   ,							split="""train"""   ,							streaming=A_				)
       __lowercase						=						ds.cast_column("""audio"""   ,							datasets.Audio(sampling_rate=1_6_0_0_0				)				)
       __lowercase						=						iter(A_				)
       __lowercase						=						next(A_				)

       __lowercase						=						AutoProcessor.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm"""				)
       __lowercase						=						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
       __lowercase						=						processor(sample["""audio"""]["""array"""]   ,							return_tensors="""pt"""				).input_values

       with torch.no_grad():
        __lowercase						=						model(A_				).logits.cpu().numpy()

       __lowercase						=						processor.decode(logits[0]   ,							output_word_offsets=A_				)

       __lowercase						=						model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
       __lowercase						=						[
           {
               """start_time""": d["""start_offset"""] * time_offset,
               """end_time""": d["""end_offset"""] * time_offset,
               """word""": d["""word"""],
           }
           for d in output["""word_offsets"""]
       ]

       __lowercase						=						"""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(A_   ,							"""word"""				)				)   ,							A_				)
       self.assertEqual(""" """.join(self.get_from_offsets(A_   ,							"""word"""				)				)   ,							output.text				)

       # output times
       __lowercase						=						torch.tensor(self.get_from_offsets(A_   ,							"""start_time"""				)				)
       __lowercase						=						torch.tensor(self.get_from_offsets(A_   ,							"""end_time"""				)				)

       # fmt: off
       __lowercase						=						torch.tensor([1.41_99, 1.65_99, 2.25_99, 3.0, 3.24, 3.59_99, 3.79_99, 4.09_99, 4.26, 4.94, 5.28, 5.65_99, 5.78, 5.94, 6.32, 6.53_99, 6.65_99]				)
       __lowercase						=						torch.tensor([1.53_99, 1.89_99, 2.9, 3.16, 3.53_99, 3.72, 4.01_99, 4.17_99, 4.76, 5.15_99, 5.55_99, 5.69_99, 5.86, 6.19_99, 6.38, 6.61_99, 6.94]				)
       # fmt: on

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




"""simple docstring"""

from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
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  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						int   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 3   ,							A_						:						int = 1   ,							A_						:						int = 1   ,							A_						:						Optional[str] = "relu"   ,							):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						nn.Convad(
           A_   ,							A_   ,							kernel_size=A_   ,							stride=A_   ,							padding=kernel_size // 2   ,							groups=A_   ,							bias=A_   ,							)
       __lowercase						=						nn.BatchNormad(A_				)
       __lowercase						=						ACTaFN[activation] if activation is not None else nn.Identity()






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						self.convolution(A_				)
       __lowercase						=						self.normalization(A_				)
       __lowercase						=						self.activation(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Dict   ,							A_						:						RegNetConfig				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						RegNetConvLayer(
           config.num_channels   ,							config.embedding_size   ,							kernel_size=3   ,							stride=2   ,							activation=config.hidden_act				)
       __lowercase						=						config.num_channels






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						pixel_values.shape[1]
       if 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."""				)
       __lowercase						=						self.embedder(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Optional[int]   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 2				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						nn.Convad(A_   ,							A_   ,							kernel_size=1   ,							stride=A_   ,							bias=A_				)
       __lowercase						=						nn.BatchNormad(A_				)






      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						Tensor				):






       '''simple docstring'''


       __lowercase						=						self.convolution(A_				)
       __lowercase						=						self.normalization(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						int   ,							A_						:						int   ,							A_						:						int				):






       '''simple docstring'''


       super().__init__()

       __lowercase						=						nn.AdaptiveAvgPoolad((1, 1)				)
       __lowercase						=						nn.Sequential(
           nn.Convad(A_   ,							A_   ,							kernel_size=1				)   ,							nn.ReLU()   ,							nn.Convad(A_   ,							A_   ,							kernel_size=1				)   ,							nn.Sigmoid()   ,							)






      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						self.pooler(A_				)
       __lowercase						=						self.attention(A_				)
       __lowercase						=						hidden_state * attention
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Any   ,							A_						:						RegNetConfig   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 1				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						in_channels != out_channels or stride != 1
       __lowercase						=						max(1   ,							out_channels // config.groups_width				)
       __lowercase						=						(
           RegNetShortCut(A_   ,							A_   ,							stride=A_				) if should_apply_shortcut else nn.Identity()
       )
       __lowercase						=						nn.Sequential(
           RegNetConvLayer(A_   ,							A_   ,							kernel_size=1   ,							activation=config.hidden_act				)   ,							RegNetConvLayer(A_   ,							A_   ,							stride=A_   ,							groups=A_   ,							activation=config.hidden_act				)   ,							RegNetConvLayer(A_   ,							A_   ,							kernel_size=1   ,							activation=A_				)   ,							)
       __lowercase						=						ACTaFN[config.hidden_act]






      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						hidden_state
       __lowercase						=						self.layer(A_				)
       __lowercase						=						self.shortcut(A_				)
       hidden_state += residual
       __lowercase						=						self.activation(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						List[str]   ,							A_						:						RegNetConfig   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 1				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						in_channels != out_channels or stride != 1
       __lowercase						=						max(1   ,							out_channels // config.groups_width				)
       __lowercase						=						(
           RegNetShortCut(A_   ,							A_   ,							stride=A_				) if should_apply_shortcut else nn.Identity()
       )
       __lowercase						=						nn.Sequential(
           RegNetConvLayer(A_   ,							A_   ,							kernel_size=1   ,							activation=config.hidden_act				)   ,							RegNetConvLayer(A_   ,							A_   ,							stride=A_   ,							groups=A_   ,							activation=config.hidden_act				)   ,							RegNetSELayer(A_   ,							reduced_channels=int(round(in_channels / 4				)				)				)   ,							RegNetConvLayer(A_   ,							A_   ,							kernel_size=1   ,							activation=A_				)   ,							)
       __lowercase						=						ACTaFN[config.hidden_act]






      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						hidden_state
       __lowercase						=						self.layer(A_				)
       __lowercase						=						self.shortcut(A_				)
       hidden_state += residual
       __lowercase						=						self.activation(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Optional[int]   ,							A_						:						RegNetConfig   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 2   ,							A_						:						int = 2   ,							):






       '''simple docstring'''


       super().__init__()

       __lowercase						=						RegNetXLayer if config.layer_type == """x""" else RegNetYLayer

       __lowercase						=						nn.Sequential(
           # downsampling is done in the first layer with stride of 2
           layer(
               A_   ,							A_   ,							A_   ,							stride=A_   ,							)   ,							*[layer(A_   ,							A_   ,							A_				) for _ in range(depth - 1				)]   ,							)






      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]   ,							A_						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						self.layers(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Tuple   ,							A_						:						RegNetConfig				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						nn.ModuleList([]				)
       # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
       self.stages.append(
           RegNetStage(
               A_   ,							config.embedding_size   ,							config.hidden_sizes[0]   ,							stride=2 if config.downsample_in_first_stage else 1   ,							depth=config.depths[0]   ,							)				)
       __lowercase						=						zip(config.hidden_sizes   ,							config.hidden_sizes[1:]				)
       for (in_channels, out_channels), depth in zip(A_   ,							config.depths[1:]				):
        self.stages.append(RegNetStage(A_   ,							A_   ,							A_   ,							depth=A_				)				)






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Tensor   ,							A_						:						bool = False   ,							A_						:						bool = True				):






       '''simple docstring'''


       __lowercase						=						() if output_hidden_states else None

       for stage_module in self.stages:
        if output_hidden_states:
         __lowercase						=						hidden_states + (hidden_state,)

        __lowercase						=						stage_module(A_				)

       if output_hidden_states:
        __lowercase						=						hidden_states + (hidden_state,)

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

       return BaseModelOutputWithNoAttention(last_hidden_state=A_   ,							hidden_states=A_				)






class  lowerCamelCase__    (		_a     ):
      a			:  List[str]           =		RegNetConfig
      a			:  int           =		"""regnet"""
      a			:  Tuple           =		"""pixel_values"""
      a			:  Dict           =		True
      def SCREAMING_SNAKE_CASE_						(		self						:						int   ,							A_						:						List[Any]				):






       '''simple docstring'''


       if isinstance(A_   ,							nn.Convad				):
        nn.init.kaiming_normal_(module.weight   ,							mode="""fan_out"""   ,							nonlinearity="""relu"""				)
       elif isinstance(A_   ,							(nn.BatchNormad, nn.GroupNorm)				):
        nn.init.constant_(module.weight   ,							1				)
        nn.init.constant_(module.bias   ,							0				)






      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						Dict   ,							A_						:						List[str]=False				):






       '''simple docstring'''


       if isinstance(A_   ,							A_				):
        __lowercase						=						value






UpperCAmelCase__					=R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

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





@add_start_docstrings(
    """The bare RegNet model outputting raw features without any specific head on top."""      , _a      , )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class  lowerCamelCase__    (		_a     ):
      def __init__(		self						:						Any   ,							A_						:						Any				):






       '''simple docstring'''


       super().__init__(A_				)
       __lowercase						=						config
       __lowercase						=						RegNetEmbeddings(A_				)
       __lowercase						=						RegNetEncoder(A_				)
       __lowercase						=						nn.AdaptiveAvgPoolad((1, 1)				)
       # Initialize weights and apply final processing
       self.post_init()






      @add_start_docstrings_to_model_forward(A_				)
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC   ,							output_type=A_   ,							config_class=_CONFIG_FOR_DOC   ,							modality="""vision"""   ,							expected_output=_EXPECTED_OUTPUT_SHAPE   ,							)
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						Tensor   ,							A_						:						Optional[bool] = None   ,							A_						:						Optional[bool] = None				):






       '''simple docstring'''


       __lowercase						=						(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       __lowercase						=						return_dict if return_dict is not None else self.config.use_return_dict

       __lowercase						=						self.embedder(A_				)

       __lowercase						=						self.encoder(
           A_   ,							output_hidden_states=A_   ,							return_dict=A_				)

       __lowercase						=						encoder_outputs[0]

       __lowercase						=						self.pooler(A_				)

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

       return BaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=A_   ,							pooler_output=A_   ,							hidden_states=encoder_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      , )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class  lowerCamelCase__    (		_a     ):
      def __init__(		self						:						int   ,							A_						:						int				):






       '''simple docstring'''


       super().__init__(A_				)
       __lowercase						=						config.num_labels
       __lowercase						=						RegNetModel(A_				)
       # classification head
       __lowercase						=						nn.Sequential(
           nn.Flatten()   ,							nn.Linear(config.hidden_sizes[-1]   ,							config.num_labels				) if config.num_labels > 0 else nn.Identity()   ,							)
       # initialize weights and apply final processing
       self.post_init()






      @add_start_docstrings_to_model_forward(A_				)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT   ,							output_type=A_   ,							config_class=_CONFIG_FOR_DOC   ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,							)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						Optional[torch.FloatTensor] = None   ,							A_						:						Optional[torch.LongTensor] = None   ,							A_						:						Optional[bool] = None   ,							A_						:						Optional[bool] = None   ,							):






       '''simple docstring'''


       __lowercase						=						return_dict if return_dict is not None else self.config.use_return_dict

       __lowercase						=						self.regnet(A_   ,							output_hidden_states=A_   ,							return_dict=A_				)

       __lowercase						=						outputs.pooler_output if return_dict else outputs[1]

       __lowercase						=						self.classifier(A_				)

       __lowercase						=						None

       if labels is not None:
        if self.config.problem_type is None:
         if self.num_labels == 1:
          __lowercase						=						"""regression"""
         elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
          __lowercase						=						"""single_label_classification"""
         else:
          __lowercase						=						"""multi_label_classification"""
        if self.config.problem_type == "regression":
         __lowercase						=						MSELoss()
         if self.num_labels == 1:
          __lowercase						=						loss_fct(logits.squeeze()   ,							labels.squeeze()				)
         else:
          __lowercase						=						loss_fct(A_   ,							A_				)
        elif self.config.problem_type == "single_label_classification":
         __lowercase						=						CrossEntropyLoss()
         __lowercase						=						loss_fct(logits.view(-1   ,							self.num_labels				)   ,							labels.view(-1				)				)
        elif self.config.problem_type == "multi_label_classification":
         __lowercase						=						BCEWithLogitsLoss()
         __lowercase						=						loss_fct(A_   ,							A_				)

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

       return ImageClassifierOutputWithNoAttention(loss=A_   ,							logits=A_   ,							hidden_states=outputs.hidden_states				)
