






import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
	import faiss

@require_faiss
class     _SCREAMING_SNAKE_CASE       (			_a      ):





	def 						_A (		self				:					int	):
		UpperCamelCase     :List[Any]          =	tempfile.mkdtemp()
		UpperCamelCase     :List[str]          =	8

		# DPR tok
		UpperCamelCase     :Union[str, Any]          =	[
		    """[UNK]""",
		    """[CLS]""",
		    """[SEP]""",
		    """[PAD]""",
		    """[MASK]""",
		    """want""",
		    """##want""",
		    """##ed""",
		    """wa""",
		    """un""",
		    """runn""",
		    """##ing""",
		    """,""",
		    """low""",
		    """lowest""",
		]
		UpperCamelCase     :List[Any]          =	os.path.join(self.tmpdirname ,  """dpr_tokenizer"""	)
		os.makedirs(__lowerCamelCase ,  exist_ok=__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	os.path.join(__lowerCamelCase ,  DPR_VOCAB_FILES_NAMES["""vocab_file"""]	)
		with open(self.vocab_file ,  """w""" ,  encoding="""utf-8"""	) as vocab_writer:
			vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]	)	)

		# BART tok
		UpperCamelCase     :List[str]          =	[
		    """l""",
		    """o""",
		    """w""",
		    """e""",
		    """r""",
		    """s""",
		    """t""",
		    """i""",
		    """d""",
		    """n""",
		    """\u0120""",
		    """\u0120l""",
		    """\u0120n""",
		    """\u0120lo""",
		    """\u0120low""",
		    """er""",
		    """\u0120lowest""",
		    """\u0120newer""",
		    """\u0120wider""",
		    """<unk>""",
		]
		UpperCamelCase     :Dict          =	dict(zip(__lowerCamelCase ,  range(len(__lowerCamelCase	)	)	)	)
		UpperCamelCase     :int          =	["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
		UpperCamelCase     :Union[str, Any]          =	{"""unk_token""": """<unk>"""}

		UpperCamelCase     :Optional[int]          =	os.path.join(self.tmpdirname ,  """bart_tokenizer"""	)
		os.makedirs(__lowerCamelCase ,  exist_ok=__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	os.path.join(__lowerCamelCase ,  BART_VOCAB_FILES_NAMES["""vocab_file"""]	)
		UpperCamelCase     :Union[str, Any]          =	os.path.join(__lowerCamelCase ,  BART_VOCAB_FILES_NAMES["""merges_file"""]	)
		with open(self.vocab_file ,  """w""" ,  encoding="""utf-8"""	) as fp:
			fp.write(json.dumps(__lowerCamelCase	) + """\n"""	)
		with open(self.merges_file ,  """w""" ,  encoding="""utf-8"""	) as fp:
			fp.write("""\n""".join(__lowerCamelCase	)	)





	def 						_A (		self				:					List[Any]	):
		return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  """dpr_tokenizer"""	)	)





	def 						_A (		self				:					Tuple	):
		return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  """dpr_tokenizer"""	)	)





	def 						_A (		self				:					Union[str, Any]	):
		return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  """bart_tokenizer"""	)	)





	def 						_A (		self				:					Tuple	):
		shutil.rmtree(self.tmpdirname	)





	def 						_A (		self				:					Tuple	):
		UpperCamelCase     :Optional[int]          =	Dataset.from_dict(
		    {
		        """id""": ["""0""", """1"""],
		        """text""": ["""foo""", """bar"""],
		        """title""": ["""Foo""", """Bar"""],
		        """embeddings""": [np.ones(self.retrieval_vector_size	), 2 * np.ones(self.retrieval_vector_size	)],
		    }	)
		dataset.add_faiss_index("""embeddings""" ,  string_factory="""Flat""" ,  metric_type=faiss.METRIC_INNER_PRODUCT	)
		return dataset





	def 						_A (		self				:					Tuple	):
		UpperCamelCase     :List[str]          =	self.get_dummy_dataset()
		UpperCamelCase     :List[str]          =	RagConfig(
		    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  )
		with patch("""transformers.models.rag.retrieval_rag.load_dataset"""	) as mock_load_dataset:
			UpperCamelCase     :List[Any]          =	dataset
			UpperCamelCase     :Dict          =	RagRetriever(
			    __lowerCamelCase ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
		return retriever





	def 						_A (		self				:					Optional[Any] ,  __lowerCamelCase				:					bool	):
		UpperCamelCase     :Union[str, Any]          =	self.get_dummy_dataset()
		UpperCamelCase     :int          =	RagConfig(
		    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  index_name="""custom""" ,  )
		if from_disk:
			UpperCamelCase     :Any          =	os.path.join(self.tmpdirname ,  """dataset"""	)
			UpperCamelCase     :List[str]          =	os.path.join(self.tmpdirname ,  """index.faiss"""	)
			dataset.get_index("""embeddings"""	).save(os.path.join(self.tmpdirname ,  """index.faiss"""	)	)
			dataset.drop_index("""embeddings"""	)
			dataset.save_to_disk(os.path.join(self.tmpdirname ,  """dataset"""	)	)
			del dataset
			UpperCamelCase     :Optional[Any]          =	RagRetriever(
			    __lowerCamelCase ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
		else:
			UpperCamelCase     :Any          =	RagRetriever(
			    __lowerCamelCase ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  index=CustomHFIndex(config.retrieval_vector_size ,  __lowerCamelCase	) ,  )
		return retriever





	def 						_A (		self				:					Any	):
		UpperCamelCase     :List[Any]          =	Dataset.from_dict(
		    {
		        """id""": ["""0""", """1"""],
		        """text""": ["""foo""", """bar"""],
		        """title""": ["""Foo""", """Bar"""],
		        """embeddings""": [np.ones(self.retrieval_vector_size + 1	), 2 * np.ones(self.retrieval_vector_size + 1	)],
		    }	)
		dataset.add_faiss_index("""embeddings""" ,  string_factory="""Flat""" ,  metric_type=faiss.METRIC_INNER_PRODUCT	)

		UpperCamelCase     :Optional[Any]          =	os.path.join(self.tmpdirname ,  """hf_bert_base.hnswSQ8_correct_phi_128.c_index"""	)
		dataset.save_faiss_index("""embeddings""" ,  index_file_name + """.index.dpr"""	)
		pickle.dump(dataset["""id"""] ,  open(index_file_name + """.index_meta.dpr""" ,  """wb"""	)	)

		UpperCamelCase     :Union[str, Any]          =	os.path.join(self.tmpdirname ,  """psgs_w100.tsv.pkl"""	)
		UpperCamelCase     :Optional[int]          =	{sample["""id"""]: [sample["""text"""], sample["""title"""]] for sample in dataset}
		pickle.dump(__lowerCamelCase ,  open(__lowerCamelCase ,  """wb"""	)	)

		UpperCamelCase     :Dict          =	RagConfig(
		    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  index_name="""legacy""" ,  index_path=self.tmpdirname ,  )
		UpperCamelCase     :List[Any]          =	RagRetriever(
		    __lowerCamelCase ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer()	)
		return retriever





	def 						_A (		self				:					Dict	):
		UpperCamelCase     :Union[str, Any]          =	1
		UpperCamelCase     :Dict          =	self.get_dummy_canonical_hf_index_retriever()
		UpperCamelCase     :Union[str, Any]          =	np.array(
		    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :List[str]          =	retriever.retrieve(__lowerCamelCase ,  n_docs=__lowerCamelCase	)
		self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)	)
		self.assertEqual(len(__lowerCamelCase	) ,  2	)
		self.assertEqual(sorted(doc_dicts[0]	) ,  ["""embeddings""", """id""", """text""", """title"""]	)
		self.assertEqual(len(doc_dicts[0]["""id"""]	) ,  __lowerCamelCase	)
		self.assertEqual(doc_dicts[0]["""id"""][0] ,  """1"""	)  # max inner product is reached with second doc
		self.assertEqual(doc_dicts[1]["""id"""][0] ,  """0"""	)  # max inner product is reached with first doc
		self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]	)





	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :List[Any]          =	self.get_dummy_canonical_hf_index_retriever()
		with tempfile.TemporaryDirectory() as tmp_dirname:
			with patch("""transformers.models.rag.retrieval_rag.load_dataset"""	) as mock_load_dataset:
				UpperCamelCase     :List[str]          =	self.get_dummy_dataset()
				retriever.save_pretrained(__lowerCamelCase	)
				UpperCamelCase     :int          =	RagRetriever.from_pretrained(__lowerCamelCase	)
				self.assertIsInstance(__lowerCamelCase ,  __lowerCamelCase	)
				UpperCamelCase     :List[Any]          =	np.array(
				    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
				UpperCamelCase     :Dict          =	retriever.retrieve(__lowerCamelCase ,  n_docs=1	)
				self.assertTrue(out is not None	)





	def 						_A (		self				:					Tuple	):
		UpperCamelCase     :int          =	1
		UpperCamelCase     :Tuple          =	self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase	)
		UpperCamelCase     :Union[str, Any]          =	np.array(
		    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Union[str, Any]          =	retriever.retrieve(__lowerCamelCase ,  n_docs=__lowerCamelCase	)
		self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)	)
		self.assertEqual(len(__lowerCamelCase	) ,  2	)
		self.assertEqual(sorted(doc_dicts[0]	) ,  ["""embeddings""", """id""", """text""", """title"""]	)
		self.assertEqual(len(doc_dicts[0]["""id"""]	) ,  __lowerCamelCase	)
		self.assertEqual(doc_dicts[0]["""id"""][0] ,  """1"""	)  # max inner product is reached with second doc
		self.assertEqual(doc_dicts[1]["""id"""][0] ,  """0"""	)  # max inner product is reached with first doc
		self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]	)





	def 						_A (		self				:					Optional[int]	):
		UpperCamelCase     :Union[str, Any]          =	self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase	)
		with tempfile.TemporaryDirectory() as tmp_dirname:
			retriever.save_pretrained(__lowerCamelCase	)
			UpperCamelCase     :str          =	RagRetriever.from_pretrained(__lowerCamelCase	)
			self.assertIsInstance(__lowerCamelCase ,  __lowerCamelCase	)
			UpperCamelCase     :List[str]          =	np.array(
			    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
			UpperCamelCase     :Optional[int]          =	retriever.retrieve(__lowerCamelCase ,  n_docs=1	)
			self.assertTrue(out is not None	)





	def 						_A (		self				:					Any	):
		UpperCamelCase     :List[str]          =	1
		UpperCamelCase     :Dict          =	self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase	)
		UpperCamelCase     :Union[str, Any]          =	np.array(
		    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :str          =	retriever.retrieve(__lowerCamelCase ,  n_docs=__lowerCamelCase	)
		self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)	)
		self.assertEqual(len(__lowerCamelCase	) ,  2	)
		self.assertEqual(sorted(doc_dicts[0]	) ,  ["""embeddings""", """id""", """text""", """title"""]	)
		self.assertEqual(len(doc_dicts[0]["""id"""]	) ,  __lowerCamelCase	)
		self.assertEqual(doc_dicts[0]["""id"""][0] ,  """1"""	)  # max inner product is reached with second doc
		self.assertEqual(doc_dicts[1]["""id"""][0] ,  """0"""	)  # max inner product is reached with first doc
		self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]	)





	def 						_A (		self				:					List[Any]	):
		UpperCamelCase     :List[Any]          =	self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase	)
		with tempfile.TemporaryDirectory() as tmp_dirname:
			retriever.save_pretrained(__lowerCamelCase	)
			UpperCamelCase     :Optional[int]          =	RagRetriever.from_pretrained(__lowerCamelCase	)
			self.assertIsInstance(__lowerCamelCase ,  __lowerCamelCase	)
			UpperCamelCase     :Tuple          =	np.array(
			    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
			UpperCamelCase     :Any          =	retriever.retrieve(__lowerCamelCase ,  n_docs=1	)
			self.assertTrue(out is not None	)





	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :Any          =	1
		UpperCamelCase     :List[Any]          =	self.get_dummy_legacy_index_retriever()
		UpperCamelCase     :Dict          =	np.array(
		    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Any          =	retriever.retrieve(__lowerCamelCase ,  n_docs=__lowerCamelCase	)
		self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)	)
		self.assertEqual(len(__lowerCamelCase	) ,  2	)
		self.assertEqual(sorted(doc_dicts[0]	) ,  ["""text""", """title"""]	)
		self.assertEqual(len(doc_dicts[0]["""text"""]	) ,  __lowerCamelCase	)
		self.assertEqual(doc_dicts[0]["""text"""][0] ,  """bar"""	)  # max inner product is reached with second doc
		self.assertEqual(doc_dicts[1]["""text"""][0] ,  """foo"""	)  # max inner product is reached with first doc
		self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]	)





	def 						_A (		self				:					Optional[int]	):
		UpperCamelCase     :Optional[int]          =	self.get_dummy_legacy_index_retriever()
		with tempfile.TemporaryDirectory() as tmp_dirname:
			retriever.save_pretrained(__lowerCamelCase	)
			UpperCamelCase     :List[Any]          =	RagRetriever.from_pretrained(__lowerCamelCase	)
			self.assertIsInstance(__lowerCamelCase ,  __lowerCamelCase	)
			UpperCamelCase     :Optional[Any]          =	np.array(
			    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
			UpperCamelCase     :Tuple          =	retriever.retrieve(__lowerCamelCase ,  n_docs=1	)
			self.assertTrue(out is not None	)





	@require_torch
	@require_tokenizers
	@require_sentencepiece
	def 						_A (		self				:					Any	):
		import torch

		UpperCamelCase     :List[Any]          =	1
		UpperCamelCase     :List[Any]          =	self.get_dummy_canonical_hf_index_retriever()
		UpperCamelCase     :int          =	[[5, 7], [10, 11]]
		UpperCamelCase     :Dict          =	np.array(
		    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
		UpperCamelCase     :Optional[int]          =	retriever(__lowerCamelCase ,  __lowerCamelCase ,  prefix=retriever.config.generator.prefix ,  n_docs=__lowerCamelCase	)
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Optional[int]          =	(
		    out["""context_input_ids"""],
		    out["""context_attention_mask"""],
		    out["""retrieved_doc_embeds"""],
		)
		self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)	)
		self.assertIsInstance(__lowerCamelCase ,  __lowerCamelCase	)
		self.assertIsInstance(__lowerCamelCase ,  __lowerCamelCase	)
		self.assertIsInstance(__lowerCamelCase ,  np.ndarray	)

		UpperCamelCase     :List[str]          =	retriever(
		    __lowerCamelCase ,  __lowerCamelCase ,  prefix=retriever.config.generator.prefix ,  n_docs=__lowerCamelCase ,  return_tensors="""pt""" ,  )
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Optional[int]          =	(  # noqa: F841
		    out["""context_input_ids"""],
		    out["""context_attention_mask"""],
		    out["""retrieved_doc_embeds"""],
		    out["""doc_ids"""],
		)
		self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)	)
		self.assertIsInstance(__lowerCamelCase ,  torch.Tensor	)
		self.assertIsInstance(__lowerCamelCase ,  torch.Tensor	)
		self.assertIsInstance(__lowerCamelCase ,  torch.Tensor	)





	@require_torch
	@require_tokenizers
	@require_sentencepiece
	def 						_A (		self				:					List[str]	):
		UpperCamelCase     :Dict          =	self.get_dpr_ctx_encoder_tokenizer()
		UpperCamelCase     :Dict          =	1
		UpperCamelCase     :List[Any]          =	self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase	)
		retriever.set_ctx_encoder_tokenizer(__lowerCamelCase	)

		UpperCamelCase     :int          =	[[5, 7], [10, 11]]
		UpperCamelCase     :List[Any]          =	np.array(
		    [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)] ,  dtype=np.floataa	)
		UpperCamelCase     :Optional[Any]          =	retriever(__lowerCamelCase ,  __lowerCamelCase ,  prefix=retriever.config.generator.prefix ,  n_docs=__lowerCamelCase	)

		self.assertEqual(
		    len(__lowerCamelCase	) ,  6	)  # check whether the retriever output consist of 6 attributes including tokenized docs
		self.assertEqual(
		    all(k in out for k in ("""tokenized_doc_ids""", """tokenized_doc_attention_mask""")	) ,  __lowerCamelCase	)  # check for doc token related keys in dictionary.



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 ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


_UpperCAmelCase		:   Optional[Any]										=       logging.get_logger(__name__)

# General docstring
_UpperCAmelCase		:   Dict										=       """ResNetConfig"""

# Base docstring
_UpperCAmelCase		:   Optional[int]										=       """microsoft/resnet-50"""
_UpperCAmelCase		:   Optional[Any]										=       [1, 2048, 7, 7]

# Image classification docstring
_UpperCAmelCase		:   Tuple										=       """microsoft/resnet-50"""
_UpperCAmelCase		:   int										=       """tiger cat"""

_UpperCAmelCase		:   Optional[Any]										=       [
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case      ,   snake_case      ,   snake_case = 3      ,   snake_case = 1      ,   snake_case = "relu"       ):
											super().__init__()
											snake_case_							=    nn.Convad(
											    snake_case      ,   snake_case      ,   kernel_size=snake_case      ,   stride=snake_case      ,   padding=kernel_size // 2      ,   bias=snake_case       )
											snake_case_							=    nn.BatchNormad(snake_case       )
											snake_case_							=    ACTaFN[activation] if activation is not None else nn.Identity()




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.convolution(snake_case       )
											snake_case_							=    self.normalization(snake_case       )
											snake_case_							=    self.activation(snake_case       )
											return hidden_state







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case       ):
											super().__init__()
											snake_case_							=    ResNetConvLayer(
											    config.num_channels      ,   config.embedding_size      ,   kernel_size=7      ,   stride=2      ,   activation=config.hidden_act       )
											snake_case_							=    nn.MaxPoolad(kernel_size=3      ,   stride=2      ,   padding=1       )
											snake_case_							=    config.num_channels




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    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.'       )
											snake_case_							=    self.embedder(snake_case       )
											snake_case_							=    self.pooler(snake_case       )
											return embedding







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case      ,   snake_case      ,   snake_case = 2       ):
											super().__init__()
											snake_case_							=    nn.Convad(snake_case      ,   snake_case      ,   kernel_size=1      ,   stride=snake_case      ,   bias=snake_case       )
											snake_case_							=    nn.BatchNormad(snake_case       )




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.convolution(snake_case       )
											snake_case_							=    self.normalization(snake_case       )
											return hidden_state







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case      ,   snake_case      ,   snake_case = 1      ,   snake_case = "relu"       ):
											super().__init__()
											snake_case_							=    in_channels != out_channels or stride != 1
											snake_case_							=    (
											    ResNetShortCut(snake_case      ,   snake_case      ,   stride=snake_case       ) if should_apply_shortcut else nn.Identity()
											)
											snake_case_							=    nn.Sequential(
											    ResNetConvLayer(snake_case      ,   snake_case      ,   stride=snake_case       )      ,   ResNetConvLayer(snake_case      ,   snake_case      ,   activation=snake_case       )      ,   )
											snake_case_							=    ACTaFN[activation]




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    hidden_state
											snake_case_							=    self.layer(snake_case       )
											snake_case_							=    self.shortcut(snake_case       )
											hidden_state += residual
											snake_case_							=    self.activation(snake_case       )
											return hidden_state







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case      ,   snake_case      ,   snake_case = 1      ,   snake_case = "relu"      ,   snake_case = 4       ):
											super().__init__()
											snake_case_							=    in_channels != out_channels or stride != 1
											snake_case_							=    out_channels // reduction
											snake_case_							=    (
											    ResNetShortCut(snake_case      ,   snake_case      ,   stride=snake_case       ) if should_apply_shortcut else nn.Identity()
											)
											snake_case_							=    nn.Sequential(
											    ResNetConvLayer(snake_case      ,   snake_case      ,   kernel_size=1       )      ,   ResNetConvLayer(snake_case      ,   snake_case      ,   stride=snake_case       )      ,   ResNetConvLayer(snake_case      ,   snake_case      ,   kernel_size=1      ,   activation=snake_case       )      ,   )
											snake_case_							=    ACTaFN[activation]




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    hidden_state
											snake_case_							=    self.layer(snake_case       )
											snake_case_							=    self.shortcut(snake_case       )
											hidden_state += residual
											snake_case_							=    self.activation(snake_case       )
											return hidden_state







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case = 2      ,   snake_case = 2      ,   ):
											super().__init__()

											snake_case_							=    ResNetBottleNeckLayer if config.layer_type == 'bottleneck' else ResNetBasicLayer

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




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    input
											for layer in self.layers:
																		snake_case_							=    layer(snake_case       )
											return hidden_state







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case       ):
											super().__init__()
											snake_case_							=    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(
											    ResNetStage(
											        snake_case      ,   config.embedding_size      ,   config.hidden_sizes[0]      ,   stride=2 if config.downsample_in_first_stage else 1      ,   depth=config.depths[0]      ,   )       )
											snake_case_							=    zip(config.hidden_sizes      ,   config.hidden_sizes[1:]       )
											for (in_channels, out_channels), depth in zip(snake_case      ,   config.depths[1:]       ):
																		self.stages.append(ResNetStage(snake_case      ,   snake_case      ,   snake_case      ,   depth=snake_case       )       )




				def     a							(					self      ,   snake_case      ,   snake_case = False      ,   snake_case = True       ):
											snake_case_							=    () if output_hidden_states else None

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

																		snake_case_							=    stage_module(snake_case       )

											if output_hidden_states:
																		snake_case_							=    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=snake_case      ,   hidden_states=snake_case      ,   )







class  lowercase				(       lowercase_			):
				__SCREAMING_SNAKE_CASE					: List[str]       			=			ResNetConfig
				__SCREAMING_SNAKE_CASE					: Any       			=			'''resnet'''
				__SCREAMING_SNAKE_CASE					: int       			=			'''pixel_values'''
				__SCREAMING_SNAKE_CASE					: Tuple       			=			True




				def     a							(					self      ,   snake_case       ):
											if isinstance(snake_case      ,   nn.Convad       ):
																		nn.init.kaiming_normal_(module.weight      ,   mode='fan_out'      ,   nonlinearity='relu'       )
											elif isinstance(snake_case      ,   (nn.BatchNormad, nn.GroupNorm)       ):
																		nn.init.constant_(module.weight      ,   1       )
																		nn.init.constant_(module.bias      ,   0       )




				def     a							(					self      ,   snake_case      ,   snake_case=False       ):
											if isinstance(snake_case      ,   snake_case       ):
																		snake_case_							=    value



_UpperCAmelCase		:   Tuple										=       R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): 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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_UpperCAmelCase		:   Optional[int]										=       R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__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 ResNet model outputting raw features without any specific head on top.'''	,  lowercase_	,  )
class  lowercase				(       lowercase_			):




				def __init__(					self      ,   snake_case       ):
											super().__init__(snake_case       )
											snake_case_							=    config
											snake_case_							=    ResNetEmbeddings(snake_case       )
											snake_case_							=    ResNetEncoder(snake_case       )
											snake_case_							=    nn.AdaptiveAvgPoolad((1, 1)       )
											# Initialize weights and apply final processing
											self.post_init()




				@add_start_docstrings_to_model_forward(snake_case       )
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC      ,   output_type=snake_case      ,   config_class=_CONFIG_FOR_DOC      ,   modality='vision'      ,   expected_output=_EXPECTED_OUTPUT_SHAPE      ,   )
				def     a							(					self      ,   snake_case      ,   snake_case = None      ,   snake_case = None       ):
											snake_case_							=    (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											snake_case_							=    return_dict if return_dict is not None else self.config.use_return_dict

											snake_case_							=    self.embedder(snake_case       )

											snake_case_							=    self.encoder(
											    snake_case      ,   output_hidden_states=snake_case      ,   return_dict=snake_case       )

											snake_case_							=    encoder_outputs[0]

											snake_case_							=    self.pooler(snake_case       )

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

											return BaseModelOutputWithPoolingAndNoAttention(
											    last_hidden_state=snake_case      ,   pooler_output=snake_case      ,   hidden_states=encoder_outputs.hidden_states      ,   )







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




				def __init__(					self      ,   snake_case       ):
											super().__init__(snake_case       )
											snake_case_							=    config.num_labels
											snake_case_							=    ResNetModel(snake_case       )
											# classification head
											snake_case_							=    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(snake_case       )
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT      ,   output_type=snake_case      ,   config_class=_CONFIG_FOR_DOC      ,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,   )
				def     a							(					self      ,   snake_case = None      ,   snake_case = None      ,   snake_case = None      ,   snake_case = None      ,   ):
											snake_case_							=    return_dict if return_dict is not None else self.config.use_return_dict

											snake_case_							=    self.resnet(snake_case      ,   output_hidden_states=snake_case      ,   return_dict=snake_case       )

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

											snake_case_							=    self.classifier(snake_case       )

											snake_case_							=    None

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

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

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







@add_start_docstrings(
    '''
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    '''	,  lowercase_	,  )
class  lowercase				(       lowercase_	,  lowercase_			):




				def __init__(					self      ,   snake_case       ):
											super().__init__(snake_case       )
											super()._init_backbone(snake_case       )

											snake_case_							=    [config.embedding_size] + config.hidden_sizes
											snake_case_							=    ResNetEmbeddings(snake_case       )
											snake_case_							=    ResNetEncoder(snake_case       )

											# initialize weights and apply final processing
											self.post_init()




				@add_start_docstrings_to_model_forward(snake_case       )
				@replace_return_docstrings(output_type=snake_case      ,   config_class=_CONFIG_FOR_DOC       )
				def     a							(					self      ,   snake_case      ,   snake_case = None      ,   snake_case = None       ):
											snake_case_							=    return_dict if return_dict is not None else self.config.use_return_dict
											snake_case_							=    (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)

											snake_case_							=    self.embedder(snake_case       )

											snake_case_							=    self.encoder(snake_case      ,   output_hidden_states=snake_case      ,   return_dict=snake_case       )

											snake_case_							=    outputs.hidden_states

											snake_case_							=    ()
											for idx, stage in enumerate(self.stage_names       ):
																		if stage in self.out_features:
																									feature_maps += (hidden_states[idx],)

											if not return_dict:
																		snake_case_							=    (feature_maps,)
																		if output_hidden_states:
																									output += (outputs.hidden_states,)
																		return output

											return BackboneOutput(
											    feature_maps=snake_case      ,   hidden_states=outputs.hidden_states if output_hidden_states else None      ,   attentions=snake_case      ,   )




