o
    ॵi                     @   s   d dl mZmZmZ d dlmZ d dlmZ d dlm	Z	 d dl
mZ d dlmZ d dlmZmZ d dlmZ d d	lmZ e Zd
gZejejejdG dd
 d
e	ZdS )    )AnyDictUnion)	Pipelines)Model)Pipeline)	PIPELINES)WavToScp)extract_pcm_from_wavload_bytes_from_url)Tasks)
get_logger'WeNetAutomaticSpeechRecognitionPipeline)module_namec                       s   e Zd ZdZ		ddeeef def fddZ			ddeee	f de
d	ed
edeeef f
ddZdeeef deeef fddZdeeef deeef fddZ  ZS )r   zASR Inference Pipeline
    Nmodelpreprocessorc                    s   t  jd||d| dS )zPuse `model` and `preprocessor` to create an asr pipeline for prediction
        )r   r   N )super__init__)selfr   r   kwargs	__class__r   k/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/pipelines/audio/asr_wenet_inference_pipeline.pyr      s   z0WeNetAutomaticSpeechRecognitionPipeline.__init__audio_inaudio_fs
recog_typeaudio_formatreturnc           	      C   s   || _ || _|| _t|trt|\| _}nt|tr#t|\| _}n|| _|d ur-|| _| j| j| jd}| 	|}| 
|d }|S )N)audior   r   
asr_result)r   r   r   
isinstancestrr   r   bytesr
   forwardpostprocess)	r   r   r   r   r   checking_audio_fsinputsoutputrstr   r   r   __call__!   s"   


z0WeNetAutomaticSpeechRecognitionPipeline.__call__r'   c                 C   s   |  ||d< |S )zDecoding
        r    )r   r   r'   r   r   r   r$   N   s   z/WeNetAutomaticSpeechRecognitionPipeline.forwardc                 C   s   |S )z process the asr results
        r   r+   r   r   r   r%   T   s   z3WeNetAutomaticSpeechRecognitionPipeline.postprocess)NN)NNN)__name__
__module____qualname____doc__r   r   r"   r	   r   r#   intr   r   r*   r$   r%   __classcell__r   r   r   r   r      s2    




"-*N)typingr   r   r   modelscope.metainfor   modelscope.modelsr   modelscope.pipelines.baser   modelscope.pipelines.builderr   modelscope.preprocessorsr	   "modelscope.utils.audio.audio_utilsr
   r   modelscope.utils.constantr   modelscope.utils.loggerr   logger__all__register_moduleauto_speech_recognitionasr_wenet_inferencer   r   r   r   r   <module>   s   