




import argparse

import torch

from transformers import (
    UniSpeechSatConfig,
    UniSpeechSatForAudioFrameClassification,
    UniSpeechSatForSequenceClassification,
    UniSpeechSatForXVector,
    WavaVecaFeatureExtractor,
    logging,
)


logging.set_verbosity_info()
UpperCamelCase__ :   Union[str, Any]						=							logging.get_logger(__name__)





def        __UpperCAmelCase						(							lowerCamelCase_ :     Optional[int]  ,							lowerCamelCase_ :     Optional[int]  ,							lowerCamelCase_ :     List[str]  )       ->  str:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : str					  =      UniSpeechSatForSequenceClassification.from_pretrained(lowerCamelCase_  ,							config=lowerCamelCase_  )
			SCREAMING_SNAKE_CASE_   : Optional[int]					  =      downstream_dict['projector.weight']
			SCREAMING_SNAKE_CASE_   : Tuple					  =      downstream_dict['projector.bias']
			SCREAMING_SNAKE_CASE_   : List[str]					  =      downstream_dict['model.post_net.linear.weight']
			SCREAMING_SNAKE_CASE_   : Dict					  =      downstream_dict['model.post_net.linear.bias']
			return model





def        __UpperCAmelCase						(							lowerCamelCase_ :     str  ,							lowerCamelCase_ :     List[str]  ,							lowerCamelCase_ :     Optional[int]  )       ->  Optional[Any]:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      UniSpeechSatForAudioFrameClassification.from_pretrained(lowerCamelCase_  ,							config=lowerCamelCase_  )
			SCREAMING_SNAKE_CASE_   : Optional[int]					  =      downstream_dict['model.linear.weight']
			SCREAMING_SNAKE_CASE_   : str					  =      downstream_dict['model.linear.bias']
			return model





def        __UpperCAmelCase						(							lowerCamelCase_ :     List[str]  ,							lowerCamelCase_ :     str  ,							lowerCamelCase_ :     Optional[Any]  )       ->  Any:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      UniSpeechSatForXVector.from_pretrained(lowerCamelCase_  ,							config=lowerCamelCase_  )
			SCREAMING_SNAKE_CASE_   : Tuple					  =      downstream_dict['connector.weight']
			SCREAMING_SNAKE_CASE_   : Any					  =      downstream_dict['connector.bias']
			for i, kernel_size in enumerate(hf_config.tdnn_kernel  ):
						SCREAMING_SNAKE_CASE_   : Any					  =      downstream_dict[
						    F'model.framelevel_feature_extractor.module.{i}.kernel.weight'
						]
						SCREAMING_SNAKE_CASE_   : Dict					  =      downstream_dict[F'model.framelevel_feature_extractor.module.{i}.kernel.bias']

			SCREAMING_SNAKE_CASE_   : Dict					  =      downstream_dict['model.utterancelevel_feature_extractor.linear1.weight']
			SCREAMING_SNAKE_CASE_   : str					  =      downstream_dict['model.utterancelevel_feature_extractor.linear1.bias']
			SCREAMING_SNAKE_CASE_   : Dict					  =      downstream_dict['model.utterancelevel_feature_extractor.linear2.weight']
			SCREAMING_SNAKE_CASE_   : List[Any]					  =      downstream_dict['model.utterancelevel_feature_extractor.linear2.bias']
			SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      downstream_dict['objective.W']
			return model





@torch.no_grad()
def        __UpperCAmelCase						(							lowerCamelCase_ :     Any  ,							lowerCamelCase_ :     int  ,							lowerCamelCase_ :     List[str]  ,							lowerCamelCase_ :     Any  )       ->  str:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      torch.load(lowerCamelCase_  ,							map_location='cpu'  )

			SCREAMING_SNAKE_CASE_   : Tuple					  =      checkpoint['Downstream']

			SCREAMING_SNAKE_CASE_   : List[Any]					  =      UniSpeechSatConfig.from_pretrained(lowerCamelCase_  )
			SCREAMING_SNAKE_CASE_   : Optional[int]					  =      WavaVecaFeatureExtractor.from_pretrained(
			    lowerCamelCase_  ,							return_attention_mask=lowerCamelCase_  ,							do_normalize=lowerCamelCase_  )

			SCREAMING_SNAKE_CASE_   : int					  =      hf_config.architectures[0]
			if arch.endswith('ForSequenceClassification'  ):
						SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      convert_classification(lowerCamelCase_  ,							lowerCamelCase_  ,							lowerCamelCase_  )
			elif arch.endswith('ForAudioFrameClassification'  ):
						SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      convert_diarization(lowerCamelCase_  ,							lowerCamelCase_  ,							lowerCamelCase_  )
			elif arch.endswith('ForXVector'  ):
						SCREAMING_SNAKE_CASE_   : Any					  =      convert_xvector(lowerCamelCase_  ,							lowerCamelCase_  ,							lowerCamelCase_  )
			else:
						raise NotImplementedError(F'S3PRL weights conversion is not supported for {arch}'  )

			if hf_config.use_weighted_layer_sum:
						SCREAMING_SNAKE_CASE_   : str					  =      checkpoint['Featurizer']['weights']

			hf_feature_extractor.save_pretrained(lowerCamelCase_  )
			hf_model.save_pretrained(lowerCamelCase_  )


if __name__ == "__main__":
						UpperCamelCase__ :   Optional[int]						=							argparse.ArgumentParser()
						parser.add_argument(
						    '''--base_model_name''', default=None, type=str, help='''Name of the huggingface pretrained base model.'''
						)
						parser.add_argument('''--config_path''', default=None, type=str, help='''Path to the huggingface classifier config.''')
						parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to the s3prl checkpoint.''')
						parser.add_argument('''--model_dump_path''', default=None, type=str, help='''Path to the final converted model.''')
						UpperCamelCase__ :   List[str]						=							parser.parse_args()
						convert_saprl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path)











import os

# Precomputes a list of the 100 first triangular numbers
lowerCamelCase      :       str    =						[int(0.5 * n * (n + 1)) for n in range(1, 101)]





def    __lowerCAmelCase   (    ):
     __lowerCAmelCase             =       os.path.dirname(os.path.realpath(__snake_case							)							)
     __lowerCAmelCase             =       os.path.join(__snake_case       ,			"words.txt"							)

     __lowerCAmelCase             =       ""
     with open(__snake_case							) as f:
          __lowerCAmelCase             =       f.readline()

     __lowerCAmelCase             =       [word.strip("\""							) for word in words.strip("\r\n"							).split(","							)]
     __lowerCAmelCase             =       [
         word
         for word in [sum(ord(__snake_case							) - 64 for x in word							) for word in words]
         if word in TRIANGULAR_NUMBERS
     ]
     return len(__snake_case							)


if __name__ == "__main__":
 print(solution())






