o 5iX$@stdZddlZddlZddlZddlZddlZddlm Z m Z ddl m Z m Z mZe GdddZGdddZdS) u MMS Forced Aligner (Romanized Text) ==================================== Language-agnostic forced alignment using Meta's MMS_FA model via torchaudio. Works on romanized/Latin-script text. Complements the native CTC aligner. MMS_FA is trained on 1000+ languages — handles romanized Indic + English equally well. Used alongside native CTC to produce dual validation scores. Usage: aligner = MMSAligner() result = aligner.align("audio.flac", "naaku konni ads gurtuntaayi") print(f"Score: {result.alignment_score}") # 0-1 N) dataclassfield)ListDictOptionalc@seZdZUdZeed<eed<eed<eed<eedZ e eed<dZ e ed <dZ e ed <d Zeed <d Zeed <defddZdS)MMSAlignmentResultz1Result of MMS forced alignment on romanized text. audio_path transcriptionalignment_score mean_log_prob)default_factory word_scoresr num_tokens num_framesaudio_duration_secprocessing_time_secreturnc CsT|j|jddt|jdt|jddd|jD|j|jt|jdt|j dd S)NdcSsg|]}t|dqS)r)round).0sr;/home/ubuntu/maya3_transcribe/src/validators/mms_aligner.py 0z.MMSAlignmentResult.to_dict.. rr r r r rrrr) rr rr r r rrrrselfrrrto_dict*s     zMMSAlignmentResult.to_dictN)__name__ __module__ __qualname____doc__str__annotations__floatrlistr rrintrrrrr!rrrrrs     rc @seZdZdZddefddZdefddZd edej fd d Z d ede e fd dZ de e de e de ede efddZd ededefddZddZdS) MMSAligneru Language-agnostic forced alignment using torchaudio MMS_FA. Tokenization: space → * (word boundary, idx 28), CTC blank = - (idx 0). Score normalization: log_prob range ~[-10, 0] mapped to [0, 1]. autodevicecCsR|dkrtjr dnd|_n||_d|_d|_d|_d|_d|_d|_ d|_ dS)Nr,cudacpuFr) torchr. is_availabler-modellabels label2idx sample_rate_setup_complete BLANK_IDXSTAR_IDX)r r-rrr__init__@s zMMSAligner.__init__rc Cs|jrdSzOtjj}||j|_|j| |_ |j |_ ddt |j D|_ |j dd|_|j dd|_td|jd t|j d |j d d|_WdStyn}ztd |WYd }~dSd }~ww)zLazy load MMS_FA model.TcSsi|]\}}||qSrr)rilrrr [rz$MMSAligner.setup..-r*r0z[MMSAligner] Ready on z (z labels, SR=)zMSAligner] Setup failed: NF)r7 torchaudio pipelinesMMS_FA get_modeltor-r3eval get_labelsr4r6 enumerater5getr8r9printlen Exception)r bundleerrrsetupPs.    zMMSAligner.setuprcCs^t|\}}t|jdkr|jdd}tj|tjdd}||j kr-t j |||j }|S)z3Load and resample audio to MMS sample rate (16kHz).)axisdtyper) sfreadrKshapemeanr1tensorfloat32 unsqueezer6rA functionalresample)r rdatasrwaveformrrr _load_audiohs  zMMSAligner._load_audiotextcCsg}|D]+}|dkr|r|d|jkr||jq||jvr3|j||jkr3||j|q|rK|d|jkrK|d|rK|d|jks=|ra|d|jkra||ra|d|jksT|S)u Convert romanized text to MMS token indices. Space → * (word boundary). Unknown chars skipped. Punctuation stripped.  r)lowerstripr9appendr5r8pop)r ratokenschrrr_text_to_tokensts   zMMSAligner._text_to_tokensrh alignmentscorescCsd}i}t|D]\}}||jkr|d7}q|||<q|d}ddt|D} d} tt||D]$\} \} } | |jkrU| t|krU| |vrQ|| }| || | d7} q1g}| D]#}|rxt|t|}tdt dd|d}||qZ|dqZ|S)z Compute per-word alignment scores from frame-level scores. Words are delimited by * (STAR_IDX) tokens. For each word, average the log-prob scores of its frames, then convert to 0-1 confidence. rrPcSsg|]}gqSrr)r_rrrrsz3MMSAligner._compute_word_scores..r?@) rHr9rangezipr8rKrfsummaxmin)r rhrkrlword_idx token_to_wordr;tokn_wordsword_log_probs token_cursor frame_idx aligned_tokscorewiresultlpsmean_lpconfrrr_compute_word_scoress0      zMMSAligner._compute_word_scoresromanized_textc Cs|js|st||dddSt}z||}|jd|j}||}|s7t||dd|t|dWSt | | |j \}}Wdn1sQwY| }tj|gtjd} tjj|| |jd\} } | } tdtd d | d } ||| d | d }t||| | |t||jd|t|d WSty}zd dl}|t||ddt|d WYd}~Sd}~ww)a Perform forced alignment of romanized text against audio. Args: audio_path: Path to audio file romanized_text: Latin-script text to align Returns: MMSAlignmentResult with alignment score rg$)rr r r rP)rr r r rrNrR)blankrnrorr)rr r r r)r7rOrtimer`rVr6rjr1no_gradr3rEr-r/rXint32rAr[ forced_alignr8rWitemrsrtrtolistrKrL traceback print_exc)r rr start_timer_ audio_durrhemissionrm token_tensorrkrlrr r rNrrrralignst         zMMSAligner.aligncCs6|jdur |`d|_d|_tjrtjdSdS)zRelease resources.NF)r3r7r1r.r2 empty_cacherrrrcleanups  zMMSAligner.cleanupN)r,)r"r#r$r%r&r:boolrOr1Tensorr`rr*rjr(rrrrrrrrr+8s"  * Or+)r%r1rA soundfilerTnumpynpr dataclassesrrtypingrrrrr+rrrrs