


"""simple docstring"""






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 							A_      (       _a     ,	_a ):
  lowerCAmelCase__      			=	'pixel_values'
  lowerCAmelCase__      			=	False
  lowerCAmelCase__      			=	TimmBackboneConfig


  def __init__(			self:						int							,__lowerCAmelCase:						List[str]							,**__lowerCAmelCase:						Tuple    ):







     '''simple docstring'''







     requires_backends(self							,"timm"    )
     super().__init__(__lowerCAmelCase    )
     _lowerCamelCase   :				Union[str, Any]							=     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."    )

     _lowerCamelCase   :				Union[str, Any]							=     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.
     _lowerCamelCase   :				int							=     config.out_indices if getattr(__lowerCAmelCase							,"out_indices"							,__lowerCAmelCase    ) is not None else (-1,)

     _lowerCamelCase   :				List[Any]							=     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.
     _lowerCamelCase   :				Tuple							=     self._backbone.return_layers
     _lowerCamelCase   :				Tuple							=     {layer["module"]: str(__lowerCAmelCase    ) for i, layer in enumerate(self._backbone.feature_info.info    )}
     super()._init_backbone(__lowerCAmelCase    )


  @classmethod
  def 							_lowercase					(			cls:						Union[str, Any]							,__lowerCAmelCase:						Tuple							,*__lowerCAmelCase:						List[str]							,**__lowerCAmelCase:						Optional[Any]    ):







     '''simple docstring'''







     requires_backends(cls							,["vision", "timm"]    )
     from ...models.timm_backbone import TimmBackboneConfig

     _lowerCamelCase   :				List[Any]							=     kwargs.pop("config"							,TimmBackboneConfig()    )

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

     _lowerCamelCase   :				str							=     kwargs.pop("num_channels"							,config.num_channels    )
     _lowerCamelCase   :				int							=     kwargs.pop("features_only"							,config.features_only    )
     _lowerCamelCase   :				Dict							=     kwargs.pop("use_pretrained_backbone"							,config.use_pretrained_backbone    )
     _lowerCamelCase   :				List[Any]							=     kwargs.pop("out_indices"							,config.out_indices    )
     _lowerCamelCase   :				Tuple							=     TimmBackboneConfig(
         backbone=__lowerCAmelCase							,num_channels=__lowerCAmelCase							,features_only=__lowerCAmelCase							,use_pretrained_backbone=__lowerCAmelCase							,out_indices=__lowerCAmelCase							,)
     return super()._from_config(__lowerCAmelCase							,**__lowerCAmelCase    )


  def 							_lowercase					(			self:						List[str]							,__lowerCAmelCase:						List[str]    ):







     '''simple docstring'''







     pass






  def 							_lowercase					(			self:						List[str]							,__lowerCAmelCase:						List[str]							,__lowerCAmelCase:						List[Any]=None							,__lowerCAmelCase:						Optional[Any]=None							,__lowerCAmelCase:						List[str]=None							,**__lowerCAmelCase:						Tuple    ):







     '''simple docstring'''







     _lowerCamelCase   :				Tuple							=     return_dict if return_dict is not None else self.config.use_return_dict
     _lowerCamelCase   :				List[Any]							=     (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     _lowerCamelCase   :				List[Any]							=     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
        _lowerCamelCase   :				Union[str, Any]							=     self._all_layers
        _lowerCamelCase   :				Tuple							=     self._backbone(__lowerCAmelCase							,**__lowerCAmelCase    )
        _lowerCamelCase   :				List[Any]							=     self._return_layers
        _lowerCamelCase   :				List[Any]							=     tuple(hidden_states[i] for i in self.out_indices    )
     else:
        _lowerCamelCase   :				Dict							=     self._backbone(__lowerCAmelCase							,**__lowerCAmelCase    )
        _lowerCamelCase   :				Union[str, Any]							=     None

     _lowerCamelCase   :				List[str]							=     tuple(__lowerCAmelCase    )
     _lowerCamelCase   :				List[Any]							=     tuple(__lowerCAmelCase    ) if hidden_states is not None else None

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

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






import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
_lowerCamelCase 					=      logging.getLogger(__name__)

if __name__ == "__main__":
	_lowerCamelCase 					=      argparse.ArgumentParser(
	    description='Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'
	)
	parser.add_argument(
	    '--data_file', type=str, default='data/dump.bert-base-uncased.pickle', help='The binarized dataset.'
	)
	parser.add_argument(
	    '--token_counts_dump', type=str, default='data/token_counts.bert-base-uncased.pickle', help='The dump file.'
	)
	parser.add_argument('--vocab_size', default=30522, type=int)
	_lowerCamelCase 					=      parser.parse_args()

	logger.info(F'''Loading data from {args.data_file}''')
	with open(args.data_file, 'rb') as fp:
		_lowerCamelCase 					=      pickle.load(fp)

	logger.info('Counting occurrences for MLM.')
	_lowerCamelCase 					=      Counter()
	for tk_ids in data:
		counter.update(tk_ids)
	_lowerCamelCase 					=      [0] * args.vocab_size
	for k, v in counter.items():
		_lowerCamelCase 					=      v

	logger.info(F'''Dump to {args.token_counts_dump}''')
	with open(args.token_counts_dump, 'wb') as handle:
		pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)