




from typing import Tuple, Union

from ...modeling_outputs import BackboneOutput
from ...modeling_utils import PreTrainedModel
from ...utils import is_timm_available, is_torch_available, requires_backends
from ...utils.backbone_utils import BackboneMixin
from .configuration_timm_backbone import TimmBackboneConfig


if is_timm_available():
		import timm


if is_torch_available():
		from torch import Tensor




class     snake_case__							(_A    ,       _A   ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_       :							Optional[Any]							 =					"pixel_values"
				SCREAMING_SNAKE_CASE_       :							Optional[Any]							 =					False
				SCREAMING_SNAKE_CASE_       :							Union[str, Any]							 =					TimmBackboneConfig





				def __init__(      self				:     str					,      __lowerCamelCase				:     Any					,      **__lowerCamelCase				:     Union[str, Any]							)					->     List[Any]:
							requires_backends(self					,      "timm"							)
							super().__init__(__lowerCamelCase							)
							a    =			config

							if config.backbone is None:
										raise ValueError("backbone is not set in the config. Please set it to a timm model name."							)

							if config.backbone not in timm.list_models():
										raise ValueError(f"""backbone {config.backbone} is not supported by timm."""							)

							if hasattr(__lowerCamelCase					,      "out_features"							) and config.out_features is not None:
										raise ValueError("out_features is not supported by TimmBackbone. Please use out_indices instead."							)

							a    =			getattr(__lowerCamelCase					,      "use_pretrained_backbone"					,      __lowerCamelCase							)
							if pretrained is None:
										raise ValueError("use_pretrained_backbone is not set in the config. Please set it to True or False."							)

							# We just take the final layer by default. This matches the default for the transformers models.
							a    =			config.out_indices if getattr(__lowerCamelCase					,      "out_indices"					,      __lowerCamelCase							) is not None else (-1,)

							a    =			timm.create_model(
							    config.backbone					,      pretrained=__lowerCamelCase					,      features_only=config.features_only					,      in_chans=config.num_channels					,      out_indices=__lowerCamelCase					,      **__lowerCamelCase					,      )
							# These are used to control the output of the model when called. If output_hidden_states is True, then
							# return_layers is modified to include all layers.
							a    =			self._backbone.return_layers
							a    =			{layer['''module''']: str(__lowerCamelCase							) for i, layer in enumerate(self._backbone.feature_info.info							)}
							super()._init_backbone(__lowerCamelCase							)





				@classmethod
				def 							__UpperCAmelCase							(      cls				:     Dict					,      __lowerCamelCase				:     Union[str, Any]					,      *__lowerCamelCase				:     Tuple					,      **__lowerCamelCase				:     str							)					->     Optional[Any]:
							requires_backends(cls					,      ["vision", "timm"]							)
							from ...models.timm_backbone import TimmBackboneConfig

							a    =			kwargs.pop("config"					,      TimmBackboneConfig()							)

							a    =			kwargs.pop("use_timm_backbone"					,      __lowerCamelCase							)
							if not use_timm:
										raise ValueError("use_timm_backbone must be True for timm backbones"							)

							a    =			kwargs.pop("num_channels"					,      config.num_channels							)
							a    =			kwargs.pop("features_only"					,      config.features_only							)
							a    =			kwargs.pop("use_pretrained_backbone"					,      config.use_pretrained_backbone							)
							a    =			kwargs.pop("out_indices"					,      config.out_indices							)
							a    =			TimmBackboneConfig(
							    backbone=__lowerCamelCase					,      num_channels=__lowerCamelCase					,      features_only=__lowerCamelCase					,      use_pretrained_backbone=__lowerCamelCase					,      out_indices=__lowerCamelCase					,      )
							return super()._from_config(__lowerCamelCase					,      **__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Tuple					,      __lowerCamelCase				:     int							)					->     int:
							pass





				def 							__UpperCAmelCase							(      self				:     Dict					,      __lowerCamelCase				:     Union[str, Any]					,      __lowerCamelCase				:     int=None					,      __lowerCamelCase				:     Optional[int]=None					,      __lowerCamelCase				:     Union[str, Any]=None					,      **__lowerCamelCase				:     Union[str, Any]							)					->     Union[BackboneOutput, Tuple[Tensor, ...]]:
							a    =			return_dict if return_dict is not None else self.config.use_return_dict
							a    =			(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							a    =			output_attentions if output_attentions is not None else self.config.output_attentions

							if output_attentions:
										raise ValueError("Cannot output attentions for timm backbones at the moment"							)

							if output_hidden_states:
										# We modify the return layers to include all the stages of the backbone
										a    =			self._all_layers
										a    =			self._backbone(__lowerCamelCase					,      **__lowerCamelCase							)
										a    =			self._return_layers
										a    =			tuple(hidden_states[i] for i in self.out_indices							)
							else:
										a    =			self._backbone(__lowerCamelCase					,      **__lowerCamelCase							)
										a    =			None

							a    =			tuple(__lowerCamelCase							)
							a    =			tuple(__lowerCamelCase							) if hidden_states is not None else None

							if not return_dict:
										a    =			(feature_maps,)
										if output_hidden_states:
													a    =			output + (hidden_states,)
										return output

							return BackboneOutput(feature_maps=__lowerCamelCase					,      hidden_states=__lowerCamelCase					,      attentions=__lowerCamelCase							)














"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__A		:       List[Any]  				=		{
    '''configuration_mega''': ['''MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MegaConfig''', '''MegaOnnxConfig'''],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:       List[str]  				=		[
						    '''MEGA_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''MegaForCausalLM''',
						    '''MegaForMaskedLM''',
						    '''MegaForMultipleChoice''',
						    '''MegaForQuestionAnswering''',
						    '''MegaForSequenceClassification''',
						    '''MegaForTokenClassification''',
						    '''MegaModel''',
						    '''MegaPreTrainedModel''',
						]

if TYPE_CHECKING:
						from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_mega import (
												    MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
												    MegaForCausalLM,
												    MegaForMaskedLM,
												    MegaForMultipleChoice,
												    MegaForQuestionAnswering,
												    MegaForSequenceClassification,
												    MegaForTokenClassification,
												    MegaModel,
												    MegaPreTrainedModel,
												)

else:
						import sys

						__A		:       Any  				=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



