o Ei@@sddlZddlZddlmZddlmZmZmZmZm Z ddl Z ddl Z ddl m Z mZmZddlmZmZmZddlmZmZmZGdd d Zed d Gd d d Zed d GdddZGdddZdS)N) dataclass)AnyDictListOptionalUnion)Qwen3ASRConfig Qwen3ASRForConditionalGenerationQwen3ASRProcessor) AutoConfig AutoModel AutoProcessor) AudioLike ensure_listnormalize_audiosc@seZdZddZdedefddZdedefdd Zdedefd d Zd ede efd dZ d ede efddZ d ede efddZ dede efddZ d ede efddZde efddZd edede efddZddZdS) Qwen3ForceAlignProcessorcCstjtjtdd}i}t|ddd}|D]}|}|s!q|d}d||<qWdn1s6wY||_d|_ dS)Nassetszkorean_dict_jieba.dictrzutf-8)encodingrg?) ospathjoindirname__file__openstripsplitko_score ko_tokenizer)self ko_dict_path ko_scoresflinewordr&Z/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft/qwen_asr/inference/qwen3_forced_aligner.py__init__&s   z!Qwen3ForceAlignProcessor.__init__chreturncCs2|dkrdSt|}|ds|drdSdS)N'TLNF) unicodedatacategory startswith)r r)catr&r&r' is_kept_char3s  z%Qwen3ForceAlignProcessor.is_kept_chartokencsdfdd|DS)Nc3s|] }|r|VqdSN)r2).0r)r r&r' <sz7Qwen3ForceAlignProcessor.clean_token..)r)r r3r&r7r' clean_token;sz$Qwen3ForceAlignProcessor.clean_tokencCst|}d|ko dknpWd|kodknpWd|ko%dknpWd|ko1dknpWd |ko=d knpWd |koId knpWd |koUdkSS)NiNii4iMiiߦii?i@ii iii)ord)r r)coder&r&r' is_cjk_char>s z$Qwen3ForceAlignProcessor.is_cjk_chartextcsdggfdd}|D]}|r||q|r)|q|q|S)Ncs2rd}|}|r|gdSdSNr4)rr9append)r3cleaned current_latinr tokensr&r' flush_latinNs   zDQwen3ForceAlignProcessor.tokenize_chinese_mixed..flush_latin)r<r?r2)r r=rDr)r&rAr'tokenize_chinese_mixedJs    z/Qwen3ForceAlignProcessor.tokenize_chinese_mixedcCs6t|j}g}|D]}||}|r||q |Sr5)nagisataggingwordsr9r?)r r=rHrCwr@r&r&r'tokenize_japanesees   z*Qwen3ForceAlignProcessor.tokenize_japanesecCs4||}g}|D]}||}|r||q |Sr5)tokenizer9r?)r rr= raw_tokensrCrIw_cleanr&r&r'tokenize_koreanns   z(Qwen3ForceAlignProcessor.tokenize_koreansegcsPggfdd}|D]}||r||q |q |S)Ncs rdgdSdSr>)r?rr&bufrCr&r' flush_buf{szFQwen3ForceAlignProcessor.split_segment_with_chinese..flush_buf)r<r?)r rOrRr)r&rPr'split_segment_with_chinesews   z3Qwen3ForceAlignProcessor.split_segment_with_chinesecCs4g}|D]}||}|r|||q|Sr5)rr9extendrS)r r=rCrOr@r&r&r'tokenize_space_langs  z,Qwen3ForceAlignProcessor.tokenize_space_langcCs|}t|}dg|}dg|}td|D]'}t|D] }||||kr=||d||kr=||d||<|||<qqt|}||}g} |} | dkr]| | || } | dksP| dg|} | D]} d| | <qh|} d}||kry| |sq|}||kr| |s|d7}||kr| |r||} | dkrd}t|dddD] }| |r| |}nqd}t||D] }| |r| |}nqt||D]$}|dur|| |<q|dur|| |<q||d||kr|n|| |<qn{d}t|dddD] }| |r | |}nqd}t||D]}| |r| |}nq|durI|durI||| d}t||D]}||||d| |<q8n%|dur\t||D]}|| |<qSn|durnt||D]}|| |<qf|}n|d7}||kszdd| DS) NrFTrcSsg|]}t|qSr&)int)r6resr&r&r' sz:Qwen3ForceAlignProcessor.fix_timestamp..)tolistlenrangemaxindexr?reversecopy)r datandpparentij max_lengthmax_idx lis_indicesidx is_normalresult anomaly_countleft_valk right_valstepr&r&r' fix_timestamps   $        "       9z&Qwen3ForceAlignProcessor.fix_timestamplanguagecCs|}|dkr||}n%|dkr0|jdur(ddlm}||jd|_||j|}n||}d|d}d|}||fS)Njapanesekoreanr) LTokenizer)scoreszz)<|audio_start|><|audio_pad|><|audio_end|>) lowerrJrsoynlp.tokenizerrwrrNrUr)r r=rt word_listrw input_textr&r&r'encode_timestamps      z)Qwen3ForceAlignProcessor.encode_timestampc CsRg}||}t|D]\}}||d}||dd}||||dq |S)NrWrr= start_timeend_time)rs enumerater?) r r{ timestamptimestamp_outputtimestamp_fixedrfr%rrr&r&r'parse_timestamps   z(Qwen3ForceAlignProcessor.parse_timestampN)__name__ __module__ __qualname__r(strboolr2r9r<rrErJrNrSrUrXrsr}rr&r&r&r'r%s    Y rT)frozenc@s*eZdZUdZeed<eed<eed<dS)ForcedAlignItema One aligned item span. Attributes: text (str): The aligned unit (cjk character or word) produced by the forced aligner processor. start_time (float): Start time in seconds. end_time (float): End time in seconds. r=rrN)rrr__doc__r__annotations__rXr&r&r&r'rs   rc@s@eZdZUdZeeed<ddZddZde defd d Z d S) ForcedAlignResultz Forced alignment output for one sample. Attributes: items (List[ForcedAlignItem]): Aligned token spans. itemscC t|jSr5)iterrr7r&r&r'__iter__+ zForcedAlignResult.__iter__cCrr5)r\rr7r&r&r'__len__.rzForcedAlignResult.__len__rkr*cCs |j|Sr5r)r rkr&r&r' __getitem__1rzForcedAlignResult.__getitem__N) rrrrrrrrrrXrr&r&r&r'r s  rc @seZdZdZdededefddZede ddfd d Z d e e e e fdefd d Zedeee efdee e e fdee e e fde efddZdee e fddZdS)Qwen3ForcedAlignera A HuggingFace-style wrapper for Qwen3-ForcedAligner model inference. This wrapper provides: - `from_pretrained()` initialization via HuggingFace AutoModel/AutoProcessor - audio input normalization (path/URL/base64/(np.ndarray, sr)) - batch and single-sample forced alignment - structured output with attribute access (`.text`, `.start_time`, `.end_time`) model processoraligner_processorcCs~||_||_||_t|dd|_|jdur/z t|j|_Wnty.td|_Ynwt |j j |_ t |j j |_ dS)Ndevicecpu)rrrgetattrrnext parameters StopIterationtorchrXconfigtimestamp_token_idfloattimestamp_segment_time)r rrrr&r&r'r(@s  zQwen3ForcedAligner.__init__pretrained_model_name_or_pathr*cKsvtdttttttttj|fi|}t|ts*t dt |dtj|dd}t }||||dS)a Load Qwen3-ForcedAligner model and initialize processors. This method: 1) Registers config/model/processor for HF auto classes. 2) Loads the model using `AutoModel.from_pretrained(...)`. 3) Initializes: - HF processor (`AutoProcessor.from_pretrained(...)`) - forced alignment text processor (`Qwen3ForceAlignProcessor()`) Args: pretrained_model_name_or_path (str): HuggingFace repo id or local directory. **kwargs: Forwarded to `AutoModel.from_pretrained(...)`. Typical examples: device_map="cuda:0", dtype=torch.bfloat16. Returns: Qwen3ForcedAligner: Initialized wrapper instance. qwen3_asrzAutoModel returned z,, expected Qwen3ASRForConditionalGeneration.T)fix_mistral_regex)rrr) r registerrr r r r from_pretrained isinstance TypeErrortyper)clsrkwargsrrrr&r&r'rTs    z"Qwen3ForcedAligner.from_pretrainedrc CsPg}|D]}|tt|ddt|ddt|dddqt|dS)Nr=r4rrrr~r)r?rrgetrr)r rritr&r&r'_to_structured_items~s z'Qwen3ForcedAligner._to_structured_itemsaudior=rtcCst|}t|}t|}t|dkrt|dkr|t|}t|t|kr.t|ksBntdt|dt|dt|g}g}t||D]\} } |j| | \} } || || qK|j||ddd} | |j j  |j j } |j j di| j}|jdd }g}t| d ||D]@\}}} |||jk}||j d }|j| |}|D]}t|d d d|d <t|dd d|d<q|||q|S)a Run forced alignment for batch or single sample. Args: audio: Audio input(s). Each item supports: - local path / https URL / base64 string - (np.ndarray, sr) All audios will be converted into mono 16k float32 arrays in [-1, 1]. text: Transcript(s) for alignment. language: Language(s) for each sample (e.g., "Chinese", "English"). Returns: List[ForcedAlignResult]: One result per sample. Each result contains `items`, and each token can be accessed via `.text`, `.start_time`, `.end_time`. rzBatch size mismatch: audio=z, text=z , language=ptT)r=rreturn_tensorspaddingrV)dim input_idsrrg@@rNr&)rrr\ ValueErrorziprr}r?rtorrdtypethinkerlogitsargmaxrrnumpyrroundr)r rr=rttexts languagesaudios word_listsaligner_input_textstlangr{aligner_input_textinputsr output_idsresultsinput_id output_idmasked_output_id timestamp_msrrr&r&r'alignsD "    zQwen3ForcedAligner.aligncCs>t|jdd}t|s dS|}|durdStdd|DS)a List supported language names for the current model. This is a thin wrapper around `self.model.get_support_languages()`. If the underlying model does not expose language constraints (returns None), this method also returns None. Returns: Optional[List[str]]: - A sorted list of supported language names (lowercased), if available. - None if the model does not provide supported languages. get_support_languagesNcSsh|]}t|qSr&)rry)r6xr&r&r' sz=Qwen3ForcedAligner.get_supported_languages..)rrcallablesorted)r fnlangsr&r&r'get_supported_languagess z*Qwen3ForcedAligner.get_supported_languagesN)rrrrr r rr( classmethodrrrrrrrrinference_moderrrrrr&r&r&r'r5s6   )  Cr)rr. dataclassesrtypingrrrrrrFr"qwen_asr.core.transformers_backendrr r transformersr r r utilsrrrrrrrr&r&r&r's  j