o i5@sdZddlZddlZddlZddlZddlmZddlmZm Z ddl Z ddl Z ddl m Z ddlmZddlmZddlmZdd lmZGd d d eZGd d d eZddZdS)zEncoder definition.N)Path)OptionalTuple)FileLock)check_argument_types) AbsEncoder) make_pad_mask) LayerNormc)seZdZdZ             d/dedededededededededededededededededed ed!ef(fd"d# Zd$efd%d&Z 'd0d(e j d)e j d*e j d$e e j e j e e j ffd+d,Zd-d.ZZS)1FairseqHubertEncodera6FairSeq Hubert encoder module, used for loading pretrained weight and finetuning Args: input_size: input dim hubert_url: url to Hubert pretrained model hubert_dir_path: directory to download the Wav2Vec2.0 pretrained model. output_size: dimension of attention normalize_before: whether to use layer_norm before the first block freeze_finetune_updates: steps that freeze all layers except output layer before tuning the whole model (nessasary to prevent overfit). dropout_rate: dropout rate activation_dropout: dropout rate in activation function attention_dropout: dropout rate in attention Hubert specific Args: Please refer to: https://github.com/pytorch/fairseq/blob/master/fairseq/models/hubert/hubert.py ./Fr皙? ?staticT@? input_size hubert_urlhubert_dir_path output_sizenormalize_beforefreeze_finetune_updates dropout_rateactivation_dropoutattention_dropout mask_length mask_probmask_selection mask_other apply_maskmask_channel_lengthmask_channel_probmask_channel_othermask_channel_selection layerdropfeature_grad_multc!stsJt||_z ddl}ddlm}Wnty/}z tdtd|d}~ww||| | | | | |||||||d}|dkr||_ t j |j t dd}t d d |Drvz d d |D}Wntyu}z|d}~wwtjd |j d ddd}t|}|jddd}t||_Wdn1swYtd|jd|jdd|jd}|j}|jdd}t||_n&t|||_ |jj|j g|dd\}|_} |d}|jj j!}t|"|_||_#t$||sz|j%j&}Wnty}ztd|d}~ww||_'||_(|j(r%t)||_*|r:||kr:t j+,t j+-|||_.nd|_.||_/|0dt 1dgdS)Nr HubertModel)Error: FairSeq is not properly installed.BPlease install FairSeq: cd ${MAIN_ROOT}/tools && make fairseq.done)dropoutrrrrrr r"r#r%r$encoder_layerdropr'dataespnetcpu) map_locationcss|]}d|vVqdS)zencoder.encoderN).0kr2r2V/home/ubuntu/.local/lib/python3.10/site-packages/espnet2/asr/encoder/hubert_encoder.py isz0FairseqHubertEncoder.__init__..cSs&i|]\}}d|vr|dd|qS)label_embs_concatzencoder.encoder.)replace)r3r4vr2r2r5 ks  z1FairseqHubertEncoder.__init__../z config.yamlrzutf-8)encodingr hubert_dict)rr@ encoder_confrF) arg_overridesstrictzQError: pretrained models should be within: 'HubertModel, Hubertctc' classes, etc. num_updatesr2)2rsuper__init__r!fairseqfairseq.models.hubert.hubertr) Exceptionprinthubert_model_pathtorchloaddeviceallitemsospathjoinsplitropenyaml safe_loadpretrained_cfgFairseqHubertPretrainEncoderencodercopydeepcopypretrained_paramsdownload_hubertcheckpoint_utilsload_model_ensemble_and_taskmodelencoder_embed_dim state_dict _output_size isinstancehubert_encoder hubert_modelencodersrr after_normnn SequentialLinear output_layerrregister_buffer LongTensor)!selfrrrrrrrrrrrrr r!r"r#r$r%r&r'rGr)erBsstate config_filefradmodelstask __class__r2r5rF/s         zFairseqHubertEncoder.__init__returncC|jSNrdrpr2r2r5rz FairseqHubertEncoder.output_sizeNxs_padilens prev_statescCst||j}|j|jk}|j|jkr|jd7_n|r3|j|jdkr3|jd7_tdn|jd7_|s@tnt |j |||j oN|j ddd}Wdn1s]wY|d}|d}~|jdd}|jdur|||}|jr||}||dfS) zForward Hubert ASR Encoder. Args: xs_pad: input tensor (B, L, D) ilens: input length (B) prev_states: Not to be used now. Returns: position embedded tensor and mask z$Start fine-tuning hubert parameters!TN) padding_maskmask features_onlyrmxr)dim)rtorNrrDlogginginforLno_grad contextlib nullcontextrhr!trainingsumrmrri)rprrrmasksft enc_outputsolensr2r2r5forwards4         zFairseqHubertEncoder.forwardcCs |jj|jddtddS)NF)rCz,Pretrained Hubert model parameters reloaded!)rhload_state_dictr]rrrr2r2r5reload_pretrained_parameterssz1FairseqHubertEncoder.reload_pretrained_parameters)r r r Frr rr rrrrTrrrrrr r})__name__ __module__ __qualname____doc__intstrboolfloatrFrrLTensorrrrr __classcell__r2r2ryr5r s      3r cseZdZdZ            d*d ed ed ed edededededededededeffdd ZddZ defddZ d+de j d e j d!e j d"e j d#e j de e j e j ee j ff d$d%Zd&d'Zd(d)ZZS),rYaFairSeq Hubert pretrain encoder module, only used for pretraining stage Args: input_size: input dim output_size: dimension of attention linear_units: dimension of feedforward layers attention_heads: the number of heads of multi head attention num_blocks: the number of encoder blocks dropout_rate: dropout rate attention_dropout_rate: dropout rate in attention hubert_dict: target dictionary for Hubert pretraining label_rate: label frame rate. -1 for sequence label sample_rate: target sample rate. use_amp: whether to use automatic mixed precision normalize_before: whether to use layer_norm before the first block r r ./dict.txtdF>rr linear_unitsattention_heads num_blocksrattention_dropout_rateactivation_dropout_rater@ label_ratecheckpoint_activations sample_rateuse_ampc s:tsJt||_| |_zddlm}ddlm}ddlm }ddlm }Wnt y@}z t dt d|d}~ww|||||||| | d }i||}||_ |D]\}}t|j |rlt|j ||q[|}| | d }|D]\}}t||rt|||qy|}||| ||j ||j|_dS) Nr) Dictionary) HubertConfigr()HubertPretrainingConfigr*r+) rbencoder_ffn_embed_dimencoder_attention_headsencoder_layers final_dimr,rrr)rr)rrErFrdrfairseq.data.dictionaryrrHrr)rrIrJcfgrPhasattrsetattr_build_dictionary dictionariesrZ)rprrrrrrrrr@rrrrkwargsrrr)rrq cfg_overideskeyvaluehubert_task_cfghubert_task_cfg_overidesrvryr2r5rFsT          z%FairseqHubertPretrainEncoder.__init__cCsVtj|r t|dgt|dgt|di||n|d|g|_dS)Nsymbolscountindices0)rQrRexistsr add_from_file add_symbolr)rp dictionaryhubert_dict_pathr2r2r5rAs     z.FairseqHubertPretrainEncoder._build_dictionaryr{cCr|r}r~rr2r2r5rLrz(FairseqHubertPretrainEncoder.output_sizeNrrys_pad ys_pad_lengthrcCsJ|t||j}|dddt|f}|j||d|gdd}|S)zForward Hubert Pretrain Encoder. Args: xs_pad: input tensor (B, L, D) ilens: input length (B) prev_states: Not to be used now. Returns: position embedded tensor and mask NTF)rr target_listr) cast_mask_embrrrNminrZ)rprrrrrrrr2r2r5rOsz$FairseqHubertPretrainEncoder.forwardcCs<|jr|jjjtjjkrtj|jj |j_dSdSdSr}) rrZmask_embdtyperLcuda HalfTensorrj Parameterhalfrr2r2r5rlsz*FairseqHubertPretrainEncoder.cast_mask_embcCs@tjt|jj|j_t d|jjj d|j dS)Nz4Hubert mask embedding re-initiallized!, z, ) rLrjrrrrbuniform_rZrrrrrrr2r2r5rpsz9FairseqHubertPretrainEncoder.reload_pretrained_parameters) rrrrrr r r rrFrFr})rrrrrrrrrFrrrLrrrrrrrr2r2ryr5rYsx     @   rYcCstj|dd|dd}tj||}t|d0tj|s2tj ||t d|nt d|dWd|SWd|S1sNwY|S) NT)exist_okr<r=z.lockzHubert model downloaded z Hubert model z already exists.) rQmakedirsrTrRrSrrrLhubdownload_url_to_filerr) model_urldir_path model_name model_pathr2r2r5r^{s   r^)rrr[rrQpathlibrtypingrrrLrVfilelockr typeguardrespnet2.asr.encoder.abs_encoderr&espnet.nets.pytorch_backend.nets_utilsr2espnet.nets.pytorch_backend.transformer.layer_normr r rYr^r2r2r2r5s&       T