U
    ,:%e                     @   s,   d dl mZmZ dgZG dd dejZdS )    )nnTensor
Wav2Letterc                       s>   e Zd ZdZdeeedd fddZeed	d
dZ  Z	S )r   au  Wav2Letter model architecture from *Wav2Letter: an End-to-End ConvNet-based Speech
    Recognition System* :cite:`collobert2016wav2letter`.

    See Also:
        * `Training example <https://github.com/pytorch/audio/tree/release/0.12/examples/pipeline_wav2letter>`__

    Args:
        num_classes (int, optional): Number of classes to be classified. (Default: ``40``)
        input_type (str, optional): Wav2Letter can use as input: ``waveform``, ``power_spectrum``
         or ``mfcc`` (Default: ``waveform``).
        num_features (int, optional): Number of input features that the network will receive (Default: ``1``).
    (   waveform   N)num_classes
input_typenum_featuresreturnc                    s  t    |dkrdn|}ttj|dddddtjddtjddd	d
ddtjddtjddd	d
ddtjddtjddd	d
ddtjddtjddd	d
ddtjddtjddd	d
ddtjddtjddd	d
ddtjddtjddd	d
ddtjddtjdddd
ddtjddtjddd
d
ddtjddtjd|d
d
ddtjdd}|dkrttj|dddddtjdd}t||| _|dkr|| _d S )Nr      0         )Zin_channelsZout_channelsZkernel_sizeZstridepaddingT)Zinplace   r      i         r      -   )Zpower_spectrumZmfcc)super__init__r   Z
SequentialZConv1dZReLUacoustic_model)selfr   r	   r
   Zacoustic_num_featuresr   Zwaveform_model	__class__ [/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/torchaudio/models/wav2letter.pyr      sD    














zWav2Letter.__init__)xr   c                 C   s   |  |}tjj|dd}|S )z
        Args:
            x (torch.Tensor): Tensor of dimension (batch_size, num_features, input_length).

        Returns:
            Tensor: Predictor tensor of dimension (batch_size, number_of_classes, input_length).
        r   )dim)r   r   Z
functionalZlog_softmax)r   r   r   r   r   forward=   s    	
zWav2Letter.forward)r   r   r   )
__name__
__module____qualname____doc__intstrr   r   r!   __classcell__r   r   r   r   r      s   'N)Ztorchr   r   __all__Moduler   r   r   r   r   <module>   s   