o
    i                     @   sh   d dl mZ d dlZd dlZd dlZd dlmZ dde	de
de
fddZddejd	ejd
e	fddZdS )    )groupbyNrateoffsetdurationc           
      C   s   t | ,}t||j }|d u r|j| }nt||j }|| |j|dd}W d    n1 s4w   Y  t|}|d urn|j|krn|j	dkrR|
d}n|j}tjj|j|d}	|	|}|jd dkrn|d}||d urv|fS |jfS )Nfloat32)dtype   r   )	orig_freqnew_freq)sf	SoundFileint
samplerateframesseekreadtorch
from_numpyndim	unsqueezeT
torchaudio
transformsResampleshapesqueeze)
wav_pathr   r   r   fstart_frameframes_to_read
audio_dataaudio_tensor	resampler r#   Z/home/ubuntu/.local/lib/python3.10/site-packages/funasr/models/fun_asr_nano/tools/utils.py
load_audio	   s$   



r%   	log_probstargetsblankc              	   C   s   g }zq|  d | d } }| jd |jd ksJ tj| ||d\}}|d t|d  }}tt	|dd dD ]0\}}||krIq@t
|}|d d }|t| }	t|||	 }
|| ||	t|
dd q@W |S    Y |S )	Nr   r   )r(   c                 S   s   | d S )Nr   r#   )itemr#   r#   r$   <lambda>(   s    zforced_align.<locals>.<lambda>)key   )token
start_timeend_timescore)r   cpur   Fforced_alignr   exptolistr   	enumeratelistlenmaxappendr)   round)r&   r'   r(   items
alignmentsscoresr-   groupstartendr0   r#   r#   r$   r3      s2   r3   )Nr   N)r   )	itertoolsr   	soundfiler   r   r   torchaudio.functional
functionalr2   r   floatr%   Tensorr3   r#   r#   r#   r$   <module>   s     