






'''simple docstring'''







import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union






class 					a__    (			enum.Enum						):
     '''simple docstring'''







     _snake_case             =  '''timm'''
     _snake_case             =  '''transformers'''







def 						__UpperCAmelCase							(     snake_case__							,			snake_case__							,			snake_case__							):






      if stage_names is None:
            raise ValueError("Stage_names must be set for transformers backbones"							)

      if out_features is not None:
            if not isinstance(UpperCamelCase__							,			(list,)							):
                  raise ValueError(f"out_features must be a list {type(UpperCamelCase__							)}"							)
            if any(feat not in stage_names for feat in out_features							):
                  raise ValueError(f"out_features must be a subset of stage_names: {stage_names} got {out_features}"							)

      if out_indices is not None:
            if not isinstance(UpperCamelCase__							,			(list, tuple)							):
                  raise ValueError(f"out_indices must be a list or tuple, got {type(UpperCamelCase__							)}"							)
            if any(idx >= len(UpperCamelCase__							) for idx in out_indices							):
                  raise ValueError("out_indices must be valid indices for stage_names {stage_names}, got {out_indices}"							)

      if out_features is not None and out_indices is not None:
            if len(UpperCamelCase__							) != len(UpperCamelCase__							):
                  raise ValueError("out_features and out_indices should have the same length if both are set"							)
            if out_features != [stage_names[idx] for idx in out_indices]:
                  raise ValueError("out_features and out_indices should correspond to the same stages if both are set"							)


def 						__UpperCAmelCase							(     snake_case__							,			snake_case__							,			snake_case__							,			):






      if out_indices is None and out_features is None:
            __A		   =      [len(UpperCamelCase__							) - 1]
            __A		   =      [stage_names[-1]]
      elif out_indices is None and out_features is not None:
            __A		   =      [stage_names.index(UpperCamelCase__							) for layer in out_features]
      elif out_features is None and out_indices is not None:
            __A		   =      [stage_names[idx] for idx in out_indices]
      return out_features, out_indices


def 						__UpperCAmelCase							(     snake_case__							,			snake_case__							,			snake_case__							,			):
      verify_out_features_out_indices(out_features=UpperCamelCase__							,			out_indices=UpperCamelCase__							,			stage_names=UpperCamelCase__							)
      __A   ,				__A		   =      _align_output_features_output_indices(
          out_features=UpperCamelCase__							,			out_indices=UpperCamelCase__							,			stage_names=UpperCamelCase__							)
      # Verify that the aligned out_features and out_indices are valid
      verify_out_features_out_indices(out_features=UpperCamelCase__							,			out_indices=UpperCamelCase__							,			stage_names=UpperCamelCase__							)
      return output_features, output_indices






class 					a__    :
     '''simple docstring'''







     _snake_case             =  None






     def     _UpperCamelCase   (self    :	str				,				lowercase__    :	Dict					)  ->       None:
           if getattr(self				,				"_backbone"				,				a_					) is None:
                 raise ValueError("self._backbone must be set before calling _init_timm_backbone"					)

           # These will diagree with the defaults for the transformers models e.g. for resnet50
           # the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
           # the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
           __A		   =      [stage["module"] for stage in self._backbone.feature_info.info]
           __A		   =      [stage["num_chs"] for stage in self._backbone.feature_info.info]
           __A		   =      self._backbone.feature_info.out_indices
           __A		   =      self._backbone.feature_info.module_name()

           # We verify the out indices and out features are valid
           verify_out_features_out_indices(
               out_features=a_				,				out_indices=a_				,				stage_names=self.stage_names					)
           __A   ,				__A		   =      out_features, out_indices






     def     _UpperCamelCase   (self    :	Optional[Any]				,				lowercase__    :	Dict					)  ->       None:
           __A		   =      getattr(a_				,				"stage_names"					)
           __A		   =      getattr(a_				,				"out_features"				,				a_					)
           __A		   =      getattr(a_				,				"out_indices"				,				a_					)

           __A		   =      stage_names
           __A   ,				__A		   =      get_aligned_output_features_output_indices(
               out_features=a_				,				out_indices=a_				,				stage_names=a_					)
           # Number of channels for each stage. This is set in the transformer backbone model init
           __A		   =      None






     def     _UpperCamelCase   (self    :	str				,				lowercase__    :	str					)  ->       None:
           __A		   =      config

           __A		   =      getattr(a_				,				"use_timm_backbone"				,				a_					)
           __A		   =      BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

           if self.backbone_type == BackboneType.TIMM:
                 self._init_timm_backbone(a_					)
           elif self.backbone_type == BackboneType.TRANSFORMERS:
                 self._init_transformers_backbone(a_					)
           else:
                 raise ValueError(f"backbone_type {self.backbone_type} not supported."					)






     @property
     def     _UpperCamelCase   (self    :	int					)  ->       Tuple:
           return self._out_features






     @out_features.setter
     def     _UpperCamelCase   (self    :	Optional[Any]				,				lowercase__    :	Any					)  ->       Any:
           __A   ,				__A		   =      get_aligned_output_features_output_indices(
               out_features=a_				,				out_indices=a_				,				stage_names=self.stage_names					)






     @property
     def     _UpperCamelCase   (self    :	Optional[Any]					)  ->       str:
           return self._out_indices






     @out_indices.setter
     def     _UpperCamelCase   (self    :	List[str]				,				lowercase__    :	List[Any]					)  ->       List[Any]:
           __A   ,				__A		   =      get_aligned_output_features_output_indices(
               out_features=a_				,				out_indices=a_				,				stage_names=self.stage_names					)






     @property
     def     _UpperCamelCase   (self    :	Tuple					)  ->       Dict:
           # the current backbones will output the number of channels for each stage
           # even if that stage is not in the out_features list.
           return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names					)}






     @property
     def     _UpperCamelCase   (self    :	Tuple					)  ->       int:
           return [self.out_feature_channels[name] for name in self.out_features]






     def     _UpperCamelCase   (self    :	Any				,				*lowercase__    :	str				,				**lowercase__    :	int					)  ->       Tuple:
           __A		   =      dict(inspect.signature(self.forward					).parameters					)
           __A		   =      {k: v for k, v in kwargs.items() if k in signature}
           return self(*a_				,				**a_					)






     def     _UpperCamelCase   (self    :	int				,				lowercase__    :	Optional[Any]				,				lowercase__    :	Union[str, Any] = None				,				lowercase__    :	List[str] = None				,				lowercase__    :	Tuple = None				,				)  ->       Any:
           raise NotImplementedError("This method should be implemented by the derived class."					)






     def     _UpperCamelCase   (self    :	Optional[Any]					)  ->       int:
           __A		   =      super().to_dict()
           __A		   =      output.pop("_out_features"					)
           __A		   =      output.pop("_out_indices"					)
           return output







class 					a__    :
     '''simple docstring'''







     @property
     def     _UpperCamelCase   (self    :	List[str]					)  ->       Optional[Any]:
           return self._out_features






     @out_features.setter
     def     _UpperCamelCase   (self    :	Tuple				,				lowercase__    :	Tuple					)  ->       Dict:
           __A   ,				__A		   =      get_aligned_output_features_output_indices(
               out_features=a_				,				out_indices=a_				,				stage_names=self.stage_names					)






     @property
     def     _UpperCamelCase   (self    :	List[str]					)  ->       Tuple:
           return self._out_indices






     @out_indices.setter
     def     _UpperCamelCase   (self    :	List[str]				,				lowercase__    :	int					)  ->       str:
           __A   ,				__A		   =      get_aligned_output_features_output_indices(
               out_features=a_				,				out_indices=a_				,				stage_names=self.stage_names					)






     def     _UpperCamelCase   (self    :	Dict					)  ->       Union[str, Any]:
           __A		   =      super().to_dict()
           __A		   =      output.pop("_out_features"					)
           __A		   =      output.pop("_out_indices"					)
           return output










'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__A      :       str      =	{'configuration_umt5': ['UMT5Config', 'UMT5OnnxConfig']}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __A      :       str      =	[
           'UMT5EncoderModel',
           'UMT5ForConditionalGeneration',
           'UMT5ForQuestionAnswering',
           'UMT5Model',
           'UMT5PreTrainedModel',
       ]

if TYPE_CHECKING:
       from .configuration_umta import UMTaConfig, UMTaOnnxConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_umta import (
                  UMTaEncoderModel,
                  UMTaForConditionalGeneration,
                  UMTaForQuestionAnswering,
                  UMTaModel,
                  UMTaPreTrainedModel,
              )
else:
       import sys

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


