


import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
      from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse('1.6'):
      snake_case__				:	Any					       =					True
      from torch.cuda.amp import autocast


snake_case__				:	Optional[Any]					       =					logging.getLogger(__name__)





@dataclass
class   SCREAMING_SNAKE_CASE_			:




       '''simple docstring'''



       _a   							=							field(
           metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}						)
       _a   							=							field(
           default=a__						,       metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}						,       )
       _a   							=							field(
           default=a__						,       metadata={"help": "Whether to freeze the feature extractor layers of the model."}						)
       _a   							=							field(
           default=a__						,       metadata={"help": "Whether to log verbose messages or not."}						,       )
       _a   							=							field(
           default=2.0						,       metadata={"help": "Maximum temperature for gumbel softmax."}						)
       _a   							=							field(
           default=0.5						,       metadata={"help": "Minimum temperature for gumbel softmax."}						)
       _a   							=							field(
           default=0.999995						,       metadata={"help": "Decay of gumbel temperature during training."}						)







def    __lowerCamelCase   (      A__						: ModelArguments						,						A__						: TrainingArguments					)     ->    List[str]:
   logging.basicConfig(
       format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s"""						,						datefmt="""%m/%d/%Y %H:%M:%S"""						,						handlers=[logging.StreamHandler(sys.stdout					)]						,						)
   lowerCamelCase_						:   int		= logging.WARNING
   if model_args.verbose_logging:
      lowerCamelCase_						:   List[Any]		= logging.DEBUG
   elif trainer_utils.is_main_process(training_args.local_rank					):
      lowerCamelCase_						:   Union[str, Any]		= logging.INFO
   logger.setLevel(A__					)





@dataclass
class   SCREAMING_SNAKE_CASE_			:




       '''simple docstring'''



       _a   							=							field(
           default=a__						,       metadata={"help": "The name of the dataset to use (via the datasets library)."}						)
       _a   							=							field(
           default=a__						,       metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}						)
       _a   							=							field(
           default="train"						,       metadata={
               "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
           }						,       )
       _a   							=							field(
           default="validation"						,       metadata={
               "help": (
                   "The name of the validation data set split to use (via the datasets library). Defaults to 'validation'"
               )
           }						,       )
       _a   							=							field(
           default="file"						,       metadata={"help": "Column in the dataset that contains speech file path. Defaults to 'file'"}						,       )
       _a   							=							field(
           default=a__						,       metadata={"help": "Overwrite the cached preprocessed datasets or not."}						)
       _a   							=							field(
           default=1						,       metadata={
               "help": "The percentage of the train set used as validation set in case there's no validation split"
           }						,       )
       _a   							=							field(
           default=a__						,       metadata={"help": "The number of processes to use for the preprocessing."}						,       )
       _a   							=							field(
           default=20.0						,       metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"}						)







@dataclass
class   SCREAMING_SNAKE_CASE_			:




       '''simple docstring'''



       _a   							=							42
       _a   							=							42
       _a   							=							"longest"
       _a   							=							None
       _a   							=							None



       def __call__(						self				:				Optional[Any]		,				__a				:				List[Dict[str, Union[List[int], torch.Tensor]]]							)					->Dict[str, torch.Tensor]:
          # reformat list to dict and set to pytorch format
          lowerCamelCase_						:   int		= self.feature_extractor.pad(
              __a		,				max_length=self.max_length		,				padding=self.padding		,				pad_to_multiple_of=self.pad_to_multiple_of		,				return_tensors="""pt"""		,				)
          lowerCamelCase_						:   Optional[Any]		= self.model._get_feat_extract_output_lengths(batch["""input_values"""].shape[-1]							)

          lowerCamelCase_						:   str		= batch["""input_values"""].shape[0]

          # make sure that no loss is computed on padded inputs
          if batch["attention_mask"] is not None:
             # compute real output lengths according to convolution formula
             lowerCamelCase_						:   Union[str, Any]		= self.model._get_feat_extract_output_lengths(batch["""attention_mask"""].sum(-1							)							).to(
                 torch.long							)

             lowerCamelCase_						:   List[Any]		= torch.zeros(
                 (batch_size, mask_indices_seq_length)		,				dtype=torch.long		,				device=batch["""input_values"""].device							)

             # these two operations makes sure that all values
             # before the output lengths indices are attended to
             lowerCamelCase_						:   Dict		= 1
             lowerCamelCase_						:   List[str]		= attention_mask.flip([-1]							).cumsum(-1							).flip([-1]							).bool()

          # sample randomly masked indices
          lowerCamelCase_						:   Optional[Any]		= _compute_mask_indices(
              (batch_size, mask_indices_seq_length)		,				self.model.config.mask_time_prob		,				self.model.config.mask_time_length		,				attention_mask=__a		,				min_masks=2		,				)

          return batch







class   SCREAMING_SNAKE_CASE_			(a__						):




       '''simple docstring'''



       def __init__(						self				:				List[Any]		,				*__a				:				str		,				__a				:				Tuple=1		,				__a				:				Dict=0		,				__a				:				Optional[int]=1.0		,				**__a				:				Tuple							)					->Any:
          super().__init__(*__a		,				**__a							)
          lowerCamelCase_						:   List[str]		= 0
          lowerCamelCase_						:   int		= max_gumbel_temp
          lowerCamelCase_						:   List[Any]		= min_gumbel_temp
          lowerCamelCase_						:   Optional[Any]		= gumbel_temp_decay




       def 				_lowerCAmelCase							(						self				:				Union[str, Any]		,				__a				:				nn.Module		,				__a				:				Dict[str, Union[torch.Tensor, Any]]							)					->torch.Tensor:
          model.train()
          lowerCamelCase_						:   Dict		= self._prepare_inputs(__a							)

          if self.use_amp:
             with autocast():
                lowerCamelCase_						:   str		= self.compute_loss(__a		,				__a							)
          else:
             lowerCamelCase_						:   Dict		= self.compute_loss(__a		,				__a							)

          if self.args.n_gpu > 1 or self.deepspeed:
             if model.module.config.ctc_loss_reduction == "mean":
                lowerCamelCase_						:   List[str]		= loss.mean()
             elif model.module.config.ctc_loss_reduction == "sum":
                lowerCamelCase_						:   List[Any]		= loss.sum() / (inputs["""mask_time_indices"""]).sum()
             else:
                raise ValueError(F'''{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']'''							)

          if self.args.gradient_accumulation_steps > 1:
             lowerCamelCase_						:   str		= loss / self.args.gradient_accumulation_steps

          if self.use_amp:
             self.scaler.scale(__a							).backward()
          elif self.use_apex:
             with amp.scale_loss(__a		,				self.optimizer							) as scaled_loss:
                scaled_loss.backward()
          elif self.deepspeed:
             self.deepspeed.backward(__a							)
          else:
             loss.backward()

          self.num_update_step += 1
          # make sure gumbel softmax temperature is decayed
          if self.args.n_gpu > 1 or self.deepspeed:
             model.module.set_gumbel_temperature(
                 max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step		,				self.min_gumbel_temp							)							)
          else:
             model.set_gumbel_temperature(
                 max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step		,				self.min_gumbel_temp							)							)

          return loss.detach()







def    __lowerCamelCase   (      )     ->    str:
   # See all possible arguments in src/transformers/training_args.py
   # or by passing the --help flag to this script.
   # We now keep distinct sets of args, for a cleaner separation of concerns.

   lowerCamelCase_						:   Tuple		= HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)					)

   lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_						:   Union[str, Any]		= parser.parse_args_into_dataclasses()
   configure_logger(A__						,						A__					)

   # Downloading and loading a dataset from the hub.
   lowerCamelCase_						:   List[str]		= load_dataset(data_args.dataset_name						,						data_args.dataset_config_name						,						cache_dir=model_args.cache_dir					)

   if "validation" not in datasets.keys():
      # make sure only "validation" and "train" keys remain"
      lowerCamelCase_						:   int		= DatasetDict()
      lowerCamelCase_						:   Any		= load_dataset(
          data_args.dataset_name						,						data_args.dataset_config_name						,						split=f'''{data_args.train_split_name}[:{data_args.validation_split_percentage}%]'''						,						cache_dir=model_args.cache_dir						,						)
      lowerCamelCase_						:   Any		= load_dataset(
          data_args.dataset_name						,						data_args.dataset_config_name						,						split=f'''{data_args.train_split_name}[{data_args.validation_split_percentage}%:]'''						,						cache_dir=model_args.cache_dir						,						)
   else:
      # make sure only "validation" and "train" keys remain"
      lowerCamelCase_						:   str		= DatasetDict()
      lowerCamelCase_						:   Optional[Any]		= load_dataset(
          data_args.dataset_name						,						data_args.dataset_config_name						,						split="""validation"""						,						cache_dir=model_args.cache_dir						,						)
      lowerCamelCase_						:   List[str]		= load_dataset(
          data_args.dataset_name						,						data_args.dataset_config_name						,						split=f'''{data_args.train_split_name}'''						,						cache_dir=model_args.cache_dir						,						)

   # only normalized-inputs-training is supported
   lowerCamelCase_						:   str		= WavaVecaFeatureExtractor.from_pretrained(
       model_args.model_name_or_path						,						cache_dir=model_args.cache_dir						,						do_normalize=A__					)

   def prepare_dataset(A__						: List[Any]					):
      # check that all files have the correct sampling rate
      lowerCamelCase_,	lowerCamelCase_						:   Any		= librosa.load(batch[data_args.speech_file_column]						,						sr=feature_extractor.sampling_rate					)
      return batch

   # load audio files into numpy arrays
   lowerCamelCase_						:   Any		= datasets.map(
       A__						,						num_proc=data_args.preprocessing_num_workers						,						remove_columns=datasets["""train"""].column_names					)

   # filter audio files that are too long
   lowerCamelCase_						:   Union[str, Any]		= vectorized_datasets.filter(
       lambda A__					: len(data["""speech"""]					) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate					)					)

   def normalize(A__						: Dict					):
      return feature_extractor(batch["""speech"""]						,						sampling_rate=feature_extractor.sampling_rate					)

   # normalize and transform to `BatchFeatures`
   lowerCamelCase_						:   Tuple		= vectorized_datasets.map(
       A__						,						batched=A__						,						num_proc=data_args.preprocessing_num_workers						,						load_from_cache_file=not data_args.overwrite_cache						,						remove_columns=vectorized_datasets["""train"""].column_names						,						)

   # pretraining is only supported for "newer" stable layer norm architecture
   # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
   lowerCamelCase_						:   Dict		= WavaVecaConfig.from_pretrained(
       model_args.model_name_or_path						,						cache_dir=model_args.cache_dir						,						gradient_checkpointing=training_args.gradient_checkpointing						,						)

   if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
      raise ValueError(
          """PreTraining is only supported for ``config.do_stable_layer_norm=True`` and"""
          """ ``config.feat_extract_norm='layer'"""					)

   lowerCamelCase_						:   Union[str, Any]		= WavaVecaForPreTraining(A__					)

   lowerCamelCase_						:   Dict		= DataCollatorForWavaVecaPretraining(model=A__						,						feature_extractor=A__					)

   lowerCamelCase_						:   Optional[Any]		= WavaVecaPreTrainer(
       model=A__						,						data_collator=A__						,						args=A__						,						train_dataset=vectorized_datasets["""train"""]						,						eval_dataset=vectorized_datasets["""validation"""]						,						tokenizer=A__						,						max_gumbel_temp=model_args.max_gumbel_temperature						,						min_gumbel_temp=model_args.min_gumbel_temperature						,						gumbel_temp_decay=model_args.gumbel_temperature_decay						,						)
   trainer.train()


if __name__ == "__main__":
      main()









import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm import create_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import BitConfig, BitForImageClassification, BitImageProcessor
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
snake_case__				:	Dict					       =					logging.get_logger(__name__)


def    __lowerCamelCase   (      A__						: Optional[Any]					)     ->    List[str]:
   lowerCamelCase_						:   int		= """huggingface/label-files"""
   lowerCamelCase_						:   Dict		= """imagenet-1k-id2label.json"""
   lowerCamelCase_						:   Optional[Any]		= json.load(open(hf_hub_download(A__						,						A__						,						repo_type="""dataset"""					)						,						"""r"""					)					)
   lowerCamelCase_						:   str		= {int(A__					): v for k, v in idalabel.items()}
   lowerCamelCase_						:   List[Any]		= {v: k for k, v in idalabel.items()}

   lowerCamelCase_						:   str		= """std_conv""" if """bit""" in model_name else False

   # note that when using BiT as backbone for ViT-hybrid checkpoints,
   # one needs to additionally set config.layer_type = "bottleneck", config.stem_type = "same",
   # config.conv_layer = "std_conv_same"
   lowerCamelCase_						:   Optional[Any]		= BitConfig(
       conv_layer=A__						,						num_labels=1000						,						idalabel=A__						,						labelaid=A__						,						)

   return config


def    __lowerCamelCase   (      A__						: str					)     ->    Any:
   if "stem.conv" in name:
      lowerCamelCase_						:   Any		= name.replace("""stem.conv"""						,						"""bit.embedder.convolution"""					)
   if "blocks" in name:
      lowerCamelCase_						:   Dict		= name.replace("""blocks"""						,						"""layers"""					)
   if "head.fc" in name:
      lowerCamelCase_						:   Optional[Any]		= name.replace("""head.fc"""						,						"""classifier.1"""					)
   if name.startswith("""norm"""					):
      lowerCamelCase_						:   int		= """bit.""" + name
   if "bit" not in name and "classifier" not in name:
      lowerCamelCase_						:   str		= """bit.encoder.""" + name

   return name


def    __lowerCamelCase   (      )     ->    List[Any]:
   lowerCamelCase_						:   List[Any]		= """http://images.cocodataset.org/val2017/000000039769.jpg"""
   lowerCamelCase_						:   Optional[Any]		= Image.open(requests.get(A__						,						stream=A__					).raw					)
   return im


@torch.no_grad()
def    __lowerCamelCase   (      A__						: List[Any]						,						A__						: List[str]						,						A__						: Tuple=False					)     ->    List[str]:
   lowerCamelCase_						:   Optional[Any]		= get_config(A__					)

   # load original model from timm
   lowerCamelCase_						:   Optional[Any]		= create_model(A__						,						pretrained=A__					)
   timm_model.eval()

   # load state_dict of original model
   lowerCamelCase_						:   Optional[int]		= timm_model.state_dict()
   for key in state_dict.copy().keys():
      lowerCamelCase_						:   int		= state_dict.pop(A__					)
      lowerCamelCase_						:   Union[str, Any]		= val.squeeze() if """head""" in key else val

   # load HuggingFace model
   lowerCamelCase_						:   Tuple		= BitForImageClassification(A__					)
   model.eval()
   model.load_state_dict(A__					)

   # create image processor
   lowerCamelCase_						:   Optional[Any]		= create_transform(**resolve_data_config({}						,						model=A__					)					)
   lowerCamelCase_						:   List[Any]		= transform.transforms

   lowerCamelCase_						:   List[Any]		= {
       """bilinear""": PILImageResampling.BILINEAR,
       """bicubic""": PILImageResampling.BICUBIC,
       """nearest""": PILImageResampling.NEAREST,
   }

   lowerCamelCase_						:   List[str]		= BitImageProcessor(
       do_resize=A__						,						size={"""shortest_edge""": timm_transforms[0].size}						,						resample=pillow_resamplings[timm_transforms[0].interpolation.value]						,						do_center_crop=A__						,						crop_size={"""height""": timm_transforms[1].size[0], """width""": timm_transforms[1].size[1]}						,						do_normalize=A__						,						image_mean=timm_transforms[-1].mean.tolist()						,						image_std=timm_transforms[-1].std.tolist()						,						)

   lowerCamelCase_						:   int		= prepare_img()
   lowerCamelCase_						:   int		= transform(A__					).unsqueeze(0					)
   lowerCamelCase_						:   List[str]		= processor(A__						,						return_tensors="""pt"""					).pixel_values

   # verify pixel values
   assert torch.allclose(A__						,						A__					)

   # verify logits
   with torch.no_grad():
      lowerCamelCase_						:   str		= model(A__					)
      lowerCamelCase_						:   int		= outputs.logits

   print("""Logits:"""						,						logits[0, :3]					)
   print("""Predicted class:"""						,						model.config.idalabel[logits.argmax(-1					).item()]					)
   lowerCamelCase_						:   List[Any]		= timm_model(A__					)
   assert timm_logits.shape == outputs.logits.shape
   assert torch.allclose(A__						,						outputs.logits						,						atol=1e-3					)
   print("""Looks ok!"""					)

   if pytorch_dump_folder_path is not None:
      Path(A__					).mkdir(exist_ok=A__					)
      print(f'''Saving model {model_name} and processor to {pytorch_dump_folder_path}'''					)
      model.save_pretrained(A__					)
      processor.save_pretrained(A__					)

   if push_to_hub:
      print(f'''Pushing model {model_name} and processor to the hub'''					)
      model.push_to_hub(f'''ybelkada/{model_name}'''					)
      processor.push_to_hub(f'''ybelkada/{model_name}'''					)


if __name__ == "__main__":
      snake_case__				:	Dict					       =					argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '--model_name',
          default='resnetv2_50x1_bitm',
          type=str,
          help='Name of the BiT timm model you\'d like to convert.',
      )
      parser.add_argument(
          '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
      )
      parser.add_argument(
          '--push_to_hub',
          action='store_true',
          help='Whether to push the model to the hub.',
      )

      snake_case__				:	int					       =					parser.parse_args()
      convert_bit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





