U ,:%eՊ@sddlmZmZddlmZmZmZddlZddlm Z dddgZ Gdd d ej j Z Gd d d ej j ZGd d d eZGdddej j eZGdddej j ZGdddej j ZGdddej j ZeeeeeeeeeeeeeeeeeeeeeedddZeedddZdS))ABCabstractmethod)ListOptionalTupleN)EmformerRNNTemformer_rnnt_baseemformer_rnnt_modelcsJeZdZdZeddfdd ZejejeejejfdddZ Z S) _TimeReductionzCoalesces frames along time dimension into a fewer number of frames with higher feature dimensionality. Args: stride (int): number of frames to merge for each output frame. N)stridereturncst||_dSN)super__init__r )selfr  __class__U/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/torchaudio/models/rnnt.pyrs z_TimeReduction.__init__inputlengthsr c Csr|j\}}}|||j}|ddd|ddf}|j|jdd}||j}|||||j}|}||fS)aForward pass. B: batch size; T: maximum input sequence length in batch; D: feature dimension of each input sequence frame. Args: input (torch.Tensor): input sequences, with shape `(B, T, D)`. lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``input``. Returns: (torch.Tensor, torch.Tensor): torch.Tensor output sequences, with shape `(B, T // stride, D * stride)` torch.Tensor output lengths, with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in output sequences. Ntrunc)Z rounding_mode)shaper divZreshape contiguous) rrrBTDZ num_framesZT_maxoutputrrrforwards  z_TimeReduction.forward) __name__ __module__ __qualname____doc__intrtorchTensorrr! __classcell__rrrrr sr cs^eZdZdZd eeeeddfdd Zej e e ej e ej e ej fdd d Z ZS) _CustomLSTMaCustom long-short-term memory (LSTM) block that applies layer normalization to internal nodes. Args: input_dim (int): input dimension. hidden_dim (int): hidden dimension. layer_norm (bool, optional): if ``True``, enables layer normalization. (Default: ``False``) layer_norm_epsilon (float, optional): value of epsilon to use in layer normalization layers (Default: 1e-5) Fh㈵>N) input_dim hidden_dim layer_normlayer_norm_epsilonr csttjj|d|| d|_tjj|d|dd|_|rjtjj||d|_tjjd||d|_ ntj |_tj |_ ||_ dS)NZbiasF)Zeps) rrr'nnLinearx2gp2g LayerNormc_normg_normZIdentityr-)rr,r-r.r/rrrrCs   z_CustomLSTM.__init__)rstater cCs|dkrD|d}tj||j|j|jd}tj||j|j|jd}n|\}}||}g}|dD]}|||}| |}| dd\} } } } | } | } | } | } | || | }| |}| | }||qdtj|dd} ||g}| |fS)aForward pass. B: batch size; T: maximum sequence length in batch; D: feature dimension of each input sequence element. Args: input (torch.Tensor): with shape `(T, B, D)`. state (List[torch.Tensor] or None): list of tensors representing internal state generated in preceding invocation of ``forward``. Returns: (torch.Tensor, List[torch.Tensor]): torch.Tensor output, with shape `(T, B, hidden_dim)`. List[torch.Tensor] list of tensors representing internal state generated in current invocation of ``forward``. N)devicedtyperr0)dim)sizer'Zzerosr-r;r<r4Zunbindr5r8chunkZsigmoidtanhr7appendstack)rrr9rhcZ gated_inputoutputsZgatesZ input_gateZ forget_gateZ cell_gateZ output_gater rrrr!Vs,      z_CustomLSTM.forward)Fr+r"r#r$r%r&boolfloatrr'r(rrrr!r)rrrrr*7s r*c @steZdZeejejeejejfdddZeejejee e ejeejeje e ejfdddZ dS) _TranscriberrcCsdSrr)rrrrrrr!sz_Transcriber.forwardrrstatesr cCsdSrr)rrrrKrrrinfersz_Transcriber.inferN) r"r#r$rr'r(rr!rrrLrrrrrIs$rIcseZdZdZddddddeeeeeeeeeeeeeeedd fd d Ze j e j e e j e j fd d dZ e j je j e j eeee j e e j e j eee j fdddZZS)_EmformerEncoderaEmformer-based recurrent neural network transducer (RNN-T) encoder (transcription network). Args: input_dim (int): feature dimension of each input sequence element. output_dim (int): feature dimension of each output sequence element. segment_length (int): length of input segment expressed as number of frames. right_context_length (int): length of right context expressed as number of frames. time_reduction_input_dim (int): dimension to scale each element in input sequences to prior to applying time reduction block. time_reduction_stride (int): factor by which to reduce length of input sequence. transformer_num_heads (int): number of attention heads in each Emformer layer. transformer_ffn_dim (int): hidden layer dimension of each Emformer layer's feedforward network. transformer_num_layers (int): number of Emformer layers to instantiate. transformer_left_context_length (int): length of left context. transformer_dropout (float, optional): transformer dropout probability. (Default: 0.0) transformer_activation (str, optional): activation function to use in each Emformer layer's feedforward network. Must be one of ("relu", "gelu", "silu"). (Default: "relu") transformer_max_memory_size (int, optional): maximum number of memory elements to use. (Default: 0) transformer_weight_init_scale_strategy (str, optional): per-layer weight initialization scaling strategy. Must be one of ("depthwise", "constant", ``None``). (Default: "depthwise") transformer_tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``) relur depthwiseF)transformer_dropouttransformer_activationtransformer_max_memory_size&transformer_weight_init_scale_strategytransformer_tanh_on_memN)r, output_dimsegment_lengthright_context_lengthtime_reduction_input_dimtime_reduction_stridetransformer_num_headstransformer_ffn_dimtransformer_num_layerstransformer_left_context_lengthrQrRrSrTrUr cs|ttjj||dd|_t||_||}t|||| ||| | | ||| ||d |_ tj|||_ tj ||_ dS)NFr1)dropout activationZleft_context_lengthrXZmax_memory_sizeZweight_init_scale_strategyZ tanh_on_mem) rrr'r2r3 input_linearr time_reductionr transformer output_linearr6r.)rr,rVrWrXrYrZr[r\r]r^rQrRrSrTrUZtransformer_input_dimrrrrs0  z_EmformerEncoder.__init__rc CsF||}|||\}}|||\}}||}||} | |fS)aForward pass for training. B: batch size; T: maximum input sequence length in batch; D: feature dimension of each input sequence frame (input_dim). Args: input (torch.Tensor): input frame sequences right-padded with right context, with shape `(B, T + right context length, D)`. lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``input``. Returns: (torch.Tensor, torch.Tensor): torch.Tensor output frame sequences, with shape `(B, T // time_reduction_stride, output_dim)`. torch.Tensor output input lengths, with shape `(B,)` and i-th element representing number of valid elements for i-th batch element in output frame sequences. )rarbrcrdr.) rrrinput_linear_outtime_reduction_outtime_reduction_lengthstransformer_outtransformer_lengthsoutput_linear_outlayer_norm_outrrrr!s    z_EmformerEncoder.forwardrJc CsN||}|||\}}|j|||\}}} ||} || } | || fS)aRForward pass for inference. B: batch size; T: maximum input sequence segment length in batch; D: feature dimension of each input sequence frame (input_dim). Args: input (torch.Tensor): input frame sequence segments right-padded with right context, with shape `(B, T + right context length, D)`. lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``input``. state (List[List[torch.Tensor]] or None): list of lists of tensors representing internal state generated in preceding invocation of ``infer``. Returns: (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]): torch.Tensor output frame sequences, with shape `(B, T // time_reduction_stride, output_dim)`. torch.Tensor output input lengths, with shape `(B,)` and i-th element representing number of valid elements for i-th batch element in output. List[List[torch.Tensor]] output states; list of lists of tensors representing internal state generated in current invocation of ``infer``. )rarbrcrLrdr.) rrrrKrerfrgrhriZtransformer_statesrjrkrrrrLs#   z_EmformerEncoder.infer)r"r#r$r%r&rHstrrGrr'r(rr!jitexportrrrLr)rrrrrMs>$,"rMc sxeZdZdZd eeeeeeeedd fdd Zd ej ej e e e ej e ej ej e e ej fd d d Z ZS) _PredictoraRecurrent neural network transducer (RNN-T) prediction network. Args: num_symbols (int): size of target token lexicon. output_dim (int): feature dimension of each output sequence element. symbol_embedding_dim (int): dimension of each target token embedding. num_lstm_layers (int): number of LSTM layers to instantiate. lstm_hidden_dim (int): output dimension of each LSTM layer. lstm_layer_norm (bool, optional): if ``True``, enables layer normalization for LSTM layers. (Default: ``False``) lstm_layer_norm_epsilon (float, optional): value of epsilon to use in LSTM layer normalization layers. (Default: 1e-5) lstm_dropout (float, optional): LSTM dropout probability. (Default: 0.0) Fr+rNN) num_symbolsrVsymbol_embedding_dimnum_lstm_layerslstm_hidden_dimlstm_layer_normlstm_layer_norm_epsilon lstm_dropoutr c sttj||_tj|_tjfddt |D|_ tjj |d|_ tj ||_tj||_||_dS)Ncs(g|] }t|dkrndqS)r)r.r/)r*).0idxrsrtrurqrr Hsz'_Predictor.__init__..)p)rrr'r2Z Embedding embeddingr6input_layer_normZ ModuleListrange lstm_layersZDropoutr_r3linearoutput_layer_normrv) rrprVrqrrrsrtrurvrryrr9s  z_Predictor.__init__)rrr9r cCs|dd}||}||}|}g}t|jD]:\} } | ||dkrJdn|| \}} ||}|| q2||} || } | ddd||fS)a#Forward pass. B: batch size; U: maximum sequence length in batch; D: feature dimension of each input sequence element. Args: input (torch.Tensor): target sequences, with shape `(B, U)` and each element mapping to a target symbol, i.e. in range `[0, num_symbols)`. lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``input``. state (List[List[torch.Tensor]] or None, optional): list of lists of tensors representing internal state generated in preceding invocation of ``forward``. (Default: ``None``) Returns: (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]): torch.Tensor output encoding sequences, with shape `(B, U, output_dim)` torch.Tensor output lengths, with shape `(B,)` and i-th element representing number of valid elements for i-th batch element in output encoding sequences. List[List[torch.Tensor]] output states; list of lists of tensors representing internal state generated in current invocation of ``forward``. r:rN) Zpermuter|r} enumeraterr_rArr)rrrr9Zinput_tbZ embedding_outZinput_layer_norm_outZlstm_outZ state_outZ layer_idxZlstmZlstm_state_outZ linear_outZoutput_layer_norm_outrrrr!Xs       z_Predictor.forward)Fr+rN)NrFrrrrro(s,#rocs\eZdZdZd eeeddfdd Zejejejeje ejejejfddd Z Z S) _Joinera@Recurrent neural network transducer (RNN-T) joint network. Args: input_dim (int): source and target input dimension. output_dim (int): output dimension. activation (str, optional): activation function to use in the joiner. Must be one of ("relu", "tanh"). (Default: "relu") rON)r,rVr`r cs\ttjj||dd|_|dkr4tj|_n$|dkrJtj|_nt d|dS)NTr1rOr@zUnsupported activation ) rrr'r2r3rZReLUr`ZTanh ValueError)rr,rVr`rrrrs z_Joiner.__init__source_encodingssource_lengthstarget_encodingstarget_lengthsr cCs:|d|d}||}||}|||fS)aForward pass for training. B: batch size; T: maximum source sequence length in batch; U: maximum target sequence length in batch; D: dimension of each source and target sequence encoding. Args: source_encodings (torch.Tensor): source encoding sequences, with shape `(B, T, D)`. source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing valid sequence length of i-th batch element in ``source_encodings``. target_encodings (torch.Tensor): target encoding sequences, with shape `(B, U, D)`. target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing valid sequence length of i-th batch element in ``target_encodings``. Returns: (torch.Tensor, torch.Tensor, torch.Tensor): torch.Tensor joint network output, with shape `(B, T, U, output_dim)`. torch.Tensor output source lengths, with shape `(B,)` and i-th element representing number of valid elements along dim 1 for i-th batch element in joint network output. torch.Tensor output target lengths, with shape `(B,)` and i-th element representing number of valid elements along dim 2 for i-th batch element in joint network output. rr:)Z unsqueezerr`r)rrrrrZjoint_encodingsZactivation_outr rrrr!s"  z_Joiner.forward)rO) r"r#r$r%r&rlrr'r(rr!r)rrrrrs  rc sdeZdZdZeeeddfdd Zdej ej ej ej e e e ej e ej ej ej e e ej fdddZ ejjej ej e e e ej e ej ej e e ej fd d d Zejjej ej e ej ej fd d dZejjej ej e e e ej e ej ej e e ej fdddZejjej ej ej ej e ej ej ej fdddZZS)ratorchaudio.models.RNNT() Recurrent neural network transducer (RNN-T) model. Note: To build the model, please use one of the factory functions. See Also: :class:`torchaudio.pipelines.RNNTBundle`: ASR pipeline with pre-trained models. Args: transcriber (torch.nn.Module): transcription network. predictor (torch.nn.Module): prediction network. joiner (torch.nn.Module): joint network. N) transcriber predictorjoinerr cs t||_||_||_dSr)rrrrr)rrrrrrrrs z RNNT.__init__)sourcesrtargetsrpredictor_stater c CsL|j||d\}}|j|||d\}}}|j||||d\}}}||||fS)aForward pass for training. B: batch size; T: maximum source sequence length in batch; U: maximum target sequence length in batch; D: feature dimension of each source sequence element. Args: sources (torch.Tensor): source frame sequences right-padded with right context, with shape `(B, T, D)`. source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``sources``. targets (torch.Tensor): target sequences, with shape `(B, U)` and each element mapping to a target symbol. target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``targets``. predictor_state (List[List[torch.Tensor]] or None, optional): list of lists of tensors representing prediction network internal state generated in preceding invocation of ``forward``. (Default: ``None``) Returns: (torch.Tensor, torch.Tensor, torch.Tensor, List[List[torch.Tensor]]): torch.Tensor joint network output, with shape `(B, max output source length, max output target length, output_dim (number of target symbols))`. torch.Tensor output source lengths, with shape `(B,)` and i-th element representing number of valid elements along dim 1 for i-th batch element in joint network output. torch.Tensor output target lengths, with shape `(B,)` and i-th element representing number of valid elements along dim 2 for i-th batch element in joint network output. List[List[torch.Tensor]] output states; list of lists of tensors representing prediction network internal state generated in current invocation of ``forward``. )rrrrr9rrrr)rrr) rrrrrrrrr rrrr!s(,   z RNNT.forward)rrr9r cCs|j|||S)aApplies transcription network to sources in streaming mode. B: batch size; T: maximum source sequence segment length in batch; D: feature dimension of each source sequence frame. Args: sources (torch.Tensor): source frame sequence segments right-padded with right context, with shape `(B, T + right context length, D)`. source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``sources``. state (List[List[torch.Tensor]] or None): list of lists of tensors representing transcription network internal state generated in preceding invocation of ``transcribe_streaming``. Returns: (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]): torch.Tensor output frame sequences, with shape `(B, T // time_reduction_stride, output_dim)`. torch.Tensor output lengths, with shape `(B,)` and i-th element representing number of valid elements for i-th batch element in output. List[List[torch.Tensor]] output states; list of lists of tensors representing transcription network internal state generated in current invocation of ``transcribe_streaming``. )rrL)rrrr9rrrtranscribe_streamings#zRNNT.transcribe_streaming)rrr cCs |||S)aApplies transcription network to sources in non-streaming mode. B: batch size; T: maximum source sequence length in batch; D: feature dimension of each source sequence frame. Args: sources (torch.Tensor): source frame sequences right-padded with right context, with shape `(B, T + right context length, D)`. source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``sources``. Returns: (torch.Tensor, torch.Tensor): torch.Tensor output frame sequences, with shape `(B, T // time_reduction_stride, output_dim)`. torch.Tensor output lengths, with shape `(B,)` and i-th element representing number of valid elements for i-th batch element in output frame sequences. )r)rrrrrr transcribeDszRNNT.transcribe)rrr9r cCs|j|||dS)aApplies prediction network to targets. B: batch size; U: maximum target sequence length in batch; D: feature dimension of each target sequence frame. Args: targets (torch.Tensor): target sequences, with shape `(B, U)` and each element mapping to a target symbol, i.e. in range `[0, num_symbols)`. target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``targets``. state (List[List[torch.Tensor]] or None): list of lists of tensors representing internal state generated in preceding invocation of ``predict``. Returns: (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]): torch.Tensor output frame sequences, with shape `(B, U, output_dim)`. torch.Tensor output lengths, with shape `(B,)` and i-th element representing number of valid elements for i-th batch element in output. List[List[torch.Tensor]] output states; list of lists of tensors representing internal state generated in current invocation of ``predict``. r)r)rrrr9rrrpredictas!z RNNT.predictrcCs"|j||||d\}}}|||fS)aApplies joint network to source and target encodings. B: batch size; T: maximum source sequence length in batch; U: maximum target sequence length in batch; D: dimension of each source and target sequence encoding. Args: source_encodings (torch.Tensor): source encoding sequences, with shape `(B, T, D)`. source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing valid sequence length of i-th batch element in ``source_encodings``. target_encodings (torch.Tensor): target encoding sequences, with shape `(B, U, D)`. target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing valid sequence length of i-th batch element in ``target_encodings``. Returns: (torch.Tensor, torch.Tensor, torch.Tensor): torch.Tensor joint network output, with shape `(B, T, U, output_dim)`. torch.Tensor output source lengths, with shape `(B,)` and i-th element representing number of valid elements along dim 1 for i-th batch element in joint network output. torch.Tensor output target lengths, with shape `(B,)` and i-th element representing number of valid elements along dim 2 for i-th batch element in joint network output. r)r)rrrrrr rrrjoins# z RNNT.join)N)r"r#r$r%rIrorrr'r(rrrr!rmrnrrrrr)rrrrrsF  C$")r, encoding_dimrprWrXrYrZr[r\r]rQrRr^rSrTrUrqrrrtrurvr cCsTt||||||||| | | | | ||d}t||||||||d}t||}t|||S)a Builds Emformer-based :class:`~torchaudio.models.RNNT`. Note: For non-streaming inference, the expectation is for `transcribe` to be called on input sequences right-concatenated with `right_context_length` frames. For streaming inference, the expectation is for `transcribe_streaming` to be called on input chunks comprising `segment_length` frames right-concatenated with `right_context_length` frames. Args: input_dim (int): dimension of input sequence frames passed to transcription network. encoding_dim (int): dimension of transcription- and prediction-network-generated encodings passed to joint network. num_symbols (int): cardinality of set of target tokens. segment_length (int): length of input segment expressed as number of frames. right_context_length (int): length of right context expressed as number of frames. time_reduction_input_dim (int): dimension to scale each element in input sequences to prior to applying time reduction block. time_reduction_stride (int): factor by which to reduce length of input sequence. transformer_num_heads (int): number of attention heads in each Emformer layer. transformer_ffn_dim (int): hidden layer dimension of each Emformer layer's feedforward network. transformer_num_layers (int): number of Emformer layers to instantiate. transformer_left_context_length (int): length of left context considered by Emformer. transformer_dropout (float): Emformer dropout probability. transformer_activation (str): activation function to use in each Emformer layer's feedforward network. Must be one of ("relu", "gelu", "silu"). transformer_max_memory_size (int): maximum number of memory elements to use. transformer_weight_init_scale_strategy (str): per-layer weight initialization scaling strategy. Must be one of ("depthwise", "constant", ``None``). transformer_tanh_on_mem (bool): if ``True``, applies tanh to memory elements. symbol_embedding_dim (int): dimension of each target token embedding. num_lstm_layers (int): number of LSTM layers to instantiate. lstm_layer_norm (bool): if ``True``, enables layer normalization for LSTM layers. lstm_layer_norm_epsilon (float): value of epsilon to use in LSTM layer normalization layers. lstm_dropout (float): LSTM dropout probability. Returns: RNNT: Emformer RNN-T model. )r,rVrWrXrYrZr[r\r]rQrRr^rSrTrU)rqrrrsrtrurv)rMrorr)r,rrprWrXrYrZr[r\r]rQrRr^rSrTrUrqrrrtrurvencoderrrrrrr s:A )rpr cCs2tdd|dddddddd d d d d dddddddS)zBuilds basic version of Emformer-based :class:`~torchaudio.models.RNNT`. Args: num_symbols (int): The size of target token lexicon. Returns: RNNT: Emformer RNN-T model. Pir0ig?ZgelurrPTigMbP?g333333?)r,rrprWrXrYrZr[r\r]rQrRr^rSrTrUrqrrrtrurv)r )rprrrr s. )abcrrtypingrrrr'Ztorchaudio.modelsr__all__r2Moduler r*rIrMrorrr&rHrlrGr r rrrrsJ  ,R`=n `