     1→# coding=utf-8
     2→# Copyright 2026 The Alibaba Qwen team.
     3→# SPDX-License-Identifier: Apache-2.0
     4→#
     5→# Licensed under the Apache License, Version 2.0 (the "License");
     6→# you may not use this file except in compliance with the License.
     7→# You may obtain a copy of the License at
     8→#
     9→#     http://www.apache.org/licenses/LICENSE-2.0
    10→#
    11→# Unless required by applicable law or agreed to in writing, software
    12→# distributed under the License is distributed on an "AS IS" BASIS,
    13→# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14→# See the License for the specific language governing permissions and
    15→# limitations under the License.
    16→from dataclasses import dataclass
    17→from typing import Any, Dict, List, Optional, Union
    18→
    19→import numpy as np
    20→import torch
    21→from qwen_asr.core.transformers_backend import (
    22→    Qwen3ASRConfig,
    23→    Qwen3ASRForConditionalGeneration,
    24→    Qwen3ASRProcessor,
    25→)
    26→from transformers import AutoConfig, AutoModel, AutoProcessor
    27→
    28→AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
    29→AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
    30→AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)
    31→
    32→from .qwen3_forced_aligner import Qwen3ForcedAligner
    33→from .utils import (
    34→    MAX_ASR_INPUT_SECONDS,
    35→    MAX_FORCE_ALIGN_INPUT_SECONDS,
    36→    SAMPLE_RATE,
    37→    SUPPORTED_LANGUAGES,
    38→    AudioChunk,
    39→    AudioLike,
    40→    chunk_list,
    41→    merge_languages,
    42→    normalize_audios,
    43→    normalize_language_name,
    44→    parse_asr_output,
    45→    split_audio_into_chunks,
    46→    validate_language,
    47→)
    48→
    49→try:
    50→    from qwen_asr.core.vllm_backend import Qwen3ASRForConditionalGeneration
    51→    from vllm import ModelRegistry
    52→    ModelRegistry.register_model("Qwen3ASRForConditionalGeneration", Qwen3ASRForConditionalGeneration)
    53→except:
    54→    pass
    55→
    56→
    57→@dataclass
    58→class ASRTranscription:
    59→    """
    60→    One transcription result.
    61→
    62→    Attributes:
    63→        language (str):
    64→            Merged language string for the sample, e.g. "Chinese" or "Chinese,English".
    65→            Empty string if unknown or silent audio.
    66→        text (str):
    67→            Transcribed text.
    68→        time_stamps (Optional[Any]):
    69→            Forced aligner output (ForcedAlignResult).
    70→            Present only when return_time_stamps=True.
    71→    """
    72→    language: str
    73→    text: str
    74→    time_stamps: Optional[Any] = None
    75→
    76→
    77→@dataclass
    78→class ASRStreamingState:
    79→    """
    80→    Streaming ASR state for one audio stream (single utterance).
    81→
    82→    Attributes:
    83→        unfixed_chunk_num (int):
    84→            For the first N chunks, do not use previous ASR result as prefix prompt (reset prefix to "").
    85→        unfixed_token_num (int):
    86→            When chunk_id >= unfixed_chunk_num, rollback the last K tokens from the accumulated text
    87→            before using it as prefix prompt, to reduce boundary jitter.
    88→        chunk_size_sec (float):
    89→            Chunk size in seconds. Audio will be fed to the model in increments of this length.
    90→        chunk_size_samples (int):
    91→            Chunk size in samples at 16kHz (derived from chunk_size_sec).
    92→        chunk_id (int):
    93→            Current chunk index (0-based).
    94→        buffer (np.ndarray):
    95→            Buffered PCM samples that are not yet consumed into a full chunk.
    96→        audio_accum (np.ndarray):
    97→            Accumulated audio from the beginning of the stream up to current time (no padding).
    98→        prompt_raw (str):
    99→            Base prompt generated by chat template (with generation prompt), without appended prefix text.
   100→        context (str):
   101→            Context string.
   102→        force_language (Optional[str]):
   103→            If provided, force output to be text-only by appending "language X<asr_text>" in prompt_raw,
   104→            consistent with non-streaming transcribe().
   105→        language (str):
   106→            Latest parsed language (updated after each chunk decode). Empty if unknown/silent.
   107→        text (str):
   108→            Latest parsed transcription text (updated after each chunk decode).
   109→        _raw_decoded (str):
   110→            Internal accumulated decoded raw text (before parse_asr_output normalization).
   111→            Used for rollback/token trimming and as prefix for prompting.
   112→    """
   113→    unfixed_chunk_num: int
   114→    unfixed_token_num: int
   115→    chunk_size_sec: float
   116→    chunk_size_samples: int
   117→
   118→    chunk_id: int
   119→    buffer: np.ndarray
   120→    audio_accum: np.ndarray
   121→
   122→    prompt_raw: str
   123→    context: str
   124→    force_language: Optional[str]
   125→
   126→    language: str
   127→    text: str
   128→    _raw_decoded: str
   129→
   130→
   131→class Qwen3ASRModel:
   132→    """
   133→    Unified inference wrapper for Qwen3-ASR with two backends:
   134→      - Transformers backend
   135→      - vLLM backend
   136→
   137→    It optionally supports time stamp output via Qwen3-ForcedAligner.
   138→
   139→    Notes:
   140→      - Each request uses a context text and exactly one audio.
   141→      - If language is provided, the prompt will force the output to be text-only by appending
   142→        "language {Language}<asr_text>" to the assistant prompt.
   143→    """
   144→
   145→    def __init__(
   146→        self,
   147→        backend: str,
   148→        model: Any,
   149→        processor: Any,
   150→        sampling_params: Optional[Any] = None,
   151→        forced_aligner: Optional[Qwen3ForcedAligner] = None,
   152→        max_inference_batch_size: int = -1,
   153→        max_new_tokens: int = 512,
   154→    ):
   155→        self.backend = backend  # "transformers" | "vllm"
   156→        self.model = model
   157→        self.processor = processor
   158→        self.sampling_params = sampling_params
   159→        self.forced_aligner = forced_aligner
   160→        self.max_inference_batch_size = int(max_inference_batch_size)
   161→        self.max_new_tokens = max_new_tokens
   162→
   163→        if backend == "transformers":
   164→            self.device = getattr(model, "device", None)
   165→            if self.device is None:
   166→                try:
   167→                    self.device = next(model.parameters()).device
   168→                except StopIteration:
   169→                    self.device = torch.device("cpu")
   170→            self.dtype = getattr(model, "dtype", torch.float32)
   171→        else:
   172→            self.device = None
   173→            self.dtype = None
   174→
   175→    @classmethod
   176→    def from_pretrained(
   177→        cls,
   178→        pretrained_model_name_or_path: str,
   179→        forced_aligner: Optional[str] = None,
   180→        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
   181→        max_inference_batch_size: int = 32,
   182→        max_new_tokens: Optional[int] = 512,
   183→        **kwargs,
   184→    ) -> "Qwen3ASRModel":
   185→        """
   186→        Initialize using Transformers backend.
   187→
   188→        Args:
   189→            pretrained_model_name_or_path:
   190→                HuggingFace repo id or local directory.
   191→            forced_aligner:
   192→                Optional forced aligner model path/repo id.
   193→            forced_aligner_kwargs:
   194→                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
   195→            max_inference_batch_size:
   196→                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
   197→            max_new_tokens:
   198→                Maximum number of tokens to generate.
   199→            **kwargs:
   200→                Forwarded to AutoModel.from_pretrained(...).
   201→
   202→        Returns:
   203→            Qwen3ASRModel
   204→        """
   205→
   206→        model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
   207→
   208→        processor = AutoProcessor.from_pretrained(pretrained_model_name_or_path, fix_mistral_regex=True)
   209→
   210→        forced_aligner_model = None
   211→        if forced_aligner is not None:
   212→            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
   213→                forced_aligner, **(forced_aligner_kwargs or {})
   214→            )
   215→
   216→        return cls(
   217→            backend="transformers",
   218→            model=model,
   219→            processor=processor,
   220→            sampling_params=None,
   221→            forced_aligner=forced_aligner_model,
   222→            max_inference_batch_size=max_inference_batch_size,
   223→            max_new_tokens=max_new_tokens,
   224→        )
   225→
   226→    @classmethod
   227→    def LLM(
   228→        cls,
   229→        model: str,
   230→        forced_aligner: Optional[str] = None,
   231→        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
   232→        max_inference_batch_size: int = -1,
   233→        max_new_tokens: Optional[int] = 4096,
   234→        **kwargs,
   235→    ) -> "Qwen3ASRModel":
   236→        """
   237→        Initialize using vLLM backend.
   238→
   239→        Import is isolated to keep vLLM optional.
   240→
   241→        Args:
   242→            model:
   243→                Model path/repo for vLLM.
   244→            forced_aligner:
   245→                Optional forced aligner model path/repo id.
   246→            forced_aligner_kwargs:
   247→                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
   248→            max_inference_batch_size:
   249→                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
   250→            max_new_tokens:
   251→                Maximum number of tokens to generate.
   252→            **kwargs:
   253→                Forwarded to vllm.LLM(...).
   254→
   255→        Returns:
   256→            Qwen3ASRModel
   257→
   258→        Raises:
   259→            ImportError: If vLLM is not installed.
   260→        """
   261→        try:
   262→            from vllm import LLM as vLLM
   263→            from vllm import SamplingParams
   264→        except Exception as e:
   265→            raise ImportError(
   266→                "vLLM is not available. Install with: pip install qwen-asr[vllm]"
   267→            ) from e
   268→
   269→        llm = vLLM(model=model, **kwargs)
   270→
   271→        processor = Qwen3ASRProcessor.from_pretrained(model, fix_mistral_regex=True)
   272→        sampling_params = SamplingParams(**({"temperature": 0.0, "max_tokens": max_new_tokens}))
   273→
   274→        forced_aligner_model = None
   275→        if forced_aligner is not None:
   276→            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
   277→                forced_aligner, **(forced_aligner_kwargs or {})
   278→            )
   279→
   280→        return cls(
   281→            backend="vllm",
   282→            model=llm,
   283→            processor=processor,
   284→            sampling_params=sampling_params,
   285→            forced_aligner=forced_aligner_model,
   286→            max_inference_batch_size=max_inference_batch_size,
   287→            max_new_tokens=None,
   288→        )
   289→
   290→    def get_supported_languages(self) -> List[str]:
   291→        """
   292→        Returns the supported language list.
   293→
   294→        Returns:
   295→            List[str]: Canonical language names.
   296→        """
   297→        return list(SUPPORTED_LANGUAGES)
   298→
   299→    @torch.no_grad()
   300→    def transcribe(
   301→        self,
   302→        audio: Union[AudioLike, List[AudioLike]],
   303→        context: Union[str, List[str]] = "",
   304→        language: Optional[Union[str, List[Optional[str]]]] = None,
   305→        return_time_stamps: bool = False,
   306→    ) -> List[ASRTranscription]:
   307→        """
   308→        Transcribe audio with optional context and optional forced alignment timestamps.
   309→
   310→        Args:
   311→            audio:
   312→                Audio input(s). Supported:
   313→                  - str: local path / URL / base64 data url
   314→                  - (np.ndarray, sr)
   315→                  - list of above
   316→            context:
   317→                Context string(s). If scalar, it will be broadcast to batch size.
   318→            language:
   319→                Optional language(s). If provided, it must be in supported languages.
   320→                If scalar, it will be broadcast to batch size.
   321→                If provided, the prompt will force output to be transcription text only.
   322→            return_time_stamps:
   323→                If True, timestamps are produced via forced aligner and merged across chunks.
   324→                This requires forced_aligner initialized.
   325→
   326→        Returns:
   327→            List[ASRTranscription]: One result per input audio.
   328→
   329→        Raises:
   330→            ValueError:
   331→                - If return_time_stamps=True but forced_aligner is not provided.
   332→                - If language is unsupported.
   333→                - If batch sizes mismatch for context/language.
   334→        """
   335→        if return_time_stamps and self.forced_aligner is None:
   336→            raise ValueError("return_time_stamps=True requires `forced_aligner` to be provided at initialization.")
   337→
   338→        wavs = normalize_audios(audio)
   339→        n = len(wavs)
   340→
   341→        ctxs = context if isinstance(context, list) else [context]
   342→        if len(ctxs) == 1 and n > 1:
   343→            ctxs = ctxs * n
   344→        if len(ctxs) != n:
   345→            raise ValueError(f"Batch size mismatch: audio={n}, context={len(ctxs)}")
   346→
   347→        langs_in: List[Optional[str]]
   348→        if language is None:
   349→            langs_in = [None] * n
   350→        else:
   351→            langs_in = language if isinstance(language, list) else [language]
   352→            if len(langs_in) == 1 and n > 1:
   353→                langs_in = langs_in * n
   354→            if len(langs_in) != n:
   355→                raise ValueError(f"Batch size mismatch: audio={n}, language={len(langs_in)}")
   356→
   357→        langs_norm: List[Optional[str]] = []
   358→        for l in langs_in:
   359→            if l is None or str(l).strip() == "":
   360→                langs_norm.append(None)
   361→            else:
   362→                ln = normalize_language_name(str(l))
   363→                validate_language(ln)
   364→                langs_norm.append(ln)
   365→
   366→        max_chunk_sec = MAX_FORCE_ALIGN_INPUT_SECONDS if return_time_stamps else MAX_ASR_INPUT_SECONDS
   367→
   368→        # chunk audios and record mapping
   369→        chunks: List[AudioChunk] = []
   370→        for i, wav in enumerate(wavs):
   371→            parts = split_audio_into_chunks(
   372→                wav=wav,
   373→                sr=SAMPLE_RATE,
   374→                max_chunk_sec=max_chunk_sec,
   375→            )
   376→            for j, (cwav, offset_sec) in enumerate(parts):
   377→                chunks.append(AudioChunk(orig_index=i, chunk_index=j, wav=cwav, sr=SAMPLE_RATE, offset_sec=offset_sec))
   378→
   379→        # run ASR on chunks
   380→        chunk_ctx: List[str] = [ctxs[c.orig_index] for c in chunks]
   381→        chunk_lang: List[Optional[str]] = [langs_norm[c.orig_index] for c in chunks]
   382→        chunk_wavs: List[np.ndarray] = [c.wav for c in chunks]
   383→        raw_outputs = self._infer_asr(chunk_ctx, chunk_wavs, chunk_lang)
   384→
   385→        # parse outputs, prepare for optional alignment
   386→        per_chunk_lang: List[str] = []
   387→        per_chunk_text: List[str] = []
   388→        for out, forced_lang in zip(raw_outputs, chunk_lang):
   389→            lang, txt = parse_asr_output(out, user_language=forced_lang)
   390→            per_chunk_lang.append(lang)
   391→            per_chunk_text.append(txt)
   392→
   393→        # forced alignment (optional)
   394→        per_chunk_align: List[Optional[Any]] = [None] * len(chunks)
   395→        if return_time_stamps:
   396→            to_align_audio = []
   397→            to_align_text = []
   398→            to_align_lang = []
   399→            to_align_idx = []
   400→
   401→            for idx, (c, txt, lang_pred) in enumerate(zip(chunks, per_chunk_text, per_chunk_lang)):
   402→                if txt.strip() == "":
   403→                    continue
   404→                to_align_audio.append((c.wav, c.sr))
   405→                to_align_text.append(txt)
   406→                to_align_lang.append(lang_pred)
   407→                to_align_idx.append(idx)
   408→
   409→            # batch align with max_inference_batch_size
   410→            aligned_results: List[Any] = []
   411→            for a_chunk, t_chunk, l_chunk in zip(
   412→                chunk_list(to_align_audio, self.max_inference_batch_size),
   413→                chunk_list(to_align_text, self.max_inference_batch_size),
   414→                chunk_list(to_align_lang, self.max_inference_batch_size),
   415→            ):
   416→                aligned_results.extend(
   417→                    self.forced_aligner.align(audio=a_chunk, text=t_chunk, language=l_chunk)
   418→                )
   419→
   420→            # offset fix
   421→            for k, idx in enumerate(to_align_idx):
   422→                c = chunks[idx]
   423→                r = aligned_results[k]
   424→                per_chunk_align[idx] = self._offset_align_result(r, c.offset_sec)
   425→
   426→        # merge chunks back to original samples
   427→        out_langs: List[List[str]] = [[] for _ in range(n)]
   428→        out_texts: List[List[str]] = [[] for _ in range(n)]
   429→        out_aligns: List[List[Any]] = [[] for _ in range(n)]
   430→
   431→        for c, lang, txt, al in zip(chunks, per_chunk_lang, per_chunk_text, per_chunk_align):
   432→            out_langs[c.orig_index].append(lang)
   433→            out_texts[c.orig_index].append(txt)
   434→            if return_time_stamps and al is not None:
   435→                out_aligns[c.orig_index].append(al)
   436→
   437→        results: List[ASRTranscription] = []
   438→        for i in range(n):
   439→            merged_text = "".join([t for t in out_texts[i] if t is not None])
   440→            merged_language = merge_languages(out_langs[i])
   441→            merged_align = None
   442→            if return_time_stamps:
   443→                merged_align = self._merge_align_results(out_aligns[i])
   444→            results.append(ASRTranscription(language=merged_language, text=merged_text, time_stamps=merged_align))
   445→
   446→        return results
   447→
   448→    def _build_messages(self, context: str, audio_payload: Any) -> List[Dict[str, Any]]:
   449→        return [
   450→            {"role": "system", "content": context or ""},
   451→            {"role": "user", "content": [{"type": "audio", "audio": audio_payload}]},
   452→        ]
   453→
   454→    def _build_text_prompt(self, context: str, force_language: Optional[str]) -> str:
   455→        """
   456→        Build the string prompt for one request.
   457→
   458→        If force_language is provided, "language X<asr_text>" is appended after the generation prompt
   459→        to request text-only output.
   460→        """
   461→        msgs = self._build_messages(context=context, audio_payload="")
   462→        base = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
   463→        if force_language:
   464→            base = base + f"language {force_language}{'<asr_text>'}"
   465→        return base
   466→
   467→    def _infer_asr(
   468→        self,
   469→        contexts: List[str],
   470→        wavs: List[np.ndarray],
   471→        languages: List[Optional[str]],
   472→    ) -> List[str]:
   473→        """
   474→        Run backend inference for chunk-level items.
   475→
   476→        Args:
   477→            contexts: List of context strings.
   478→            wavs: List of mono waveforms (np.ndarray).
   479→            languages: List of forced languages or None.
   480→
   481→        Returns:
   482→            List[str]: Raw decoded strings (one per chunk).
   483→        """
   484→        if self.backend == "transformers":
   485→            return self._infer_asr_transformers(contexts, wavs, languages)
   486→        if self.backend == "vllm":
   487→            return self._infer_asr_vllm(contexts, wavs, languages)
   488→        raise RuntimeError(f"Unknown backend: {self.backend}")
   489→
   490→    def _infer_asr_transformers(
   491→        self,
   492→        contexts: List[str],
   493→        wavs: List[np.ndarray],
   494→        languages: List[Optional[str]],
   495→    ) -> List[str]:
   496→        outs: List[str] = []
   497→
   498→        texts = [self._build_text_prompt(context=c, force_language=fl) for c, fl in zip(contexts, languages)]
   499→
   500→        batch_size = self.max_inference_batch_size
   501→        if batch_size is None or batch_size < 0:
   502→            batch_size = len(texts)
   503→
   504→        for i in range(0, len(texts), batch_size):
   505→            sub_text = texts[i : i + batch_size]
   506→            sub_wavs = wavs[i : i + batch_size]
   507→            inputs = self.processor(text=sub_text, audio=sub_wavs, return_tensors="pt", padding=True)
   508→            inputs = inputs.to(self.model.device).to(self.model.dtype)
   509→
   510→            text_ids = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens)
   511→
   512→            decoded = self.processor.batch_decode(
   513→                text_ids.sequences[:, inputs["input_ids"].shape[1]:],
   514→                skip_special_tokens=True,
   515→                clean_up_tokenization_spaces=False,
   516→            )
   517→            outs.extend(list(decoded))
   518→
   519→        return outs
   520→
   521→    def _infer_asr_vllm(
   522→        self,
   523→        contexts: List[str],
   524→        wavs: List[np.ndarray],
   525→        languages: List[Optional[str]],
   526→    ) -> List[str]:
   527→        inputs: List[Dict[str, Any]] = []
   528→        for c, w, fl in zip(contexts, wavs, languages):
   529→            prompt = self._build_text_prompt(context=c, force_language=fl)
   530→            inputs.append({"prompt": prompt, "multi_modal_data": {"audio": [w]}})
   531→
   532→        outs: List[str] = []
   533→        for batch in chunk_list(inputs, self.max_inference_batch_size):
   534→            outputs = self.model.generate(batch, sampling_params=self.sampling_params, use_tqdm=False)
   535→            for o in outputs:
   536→                outs.append(o.outputs[0].text)
   537→        return outs
   538→
   539→    def _offset_align_result(self, result: Any, offset_sec: float) -> Any:
   540→        """
   541→        Apply time offset to a ForcedAlignResult-like object.
   542→
   543→        This function assumes:
   544→          - result has attribute `.items` which is a list of items with start_time/end_time in seconds.
   545→          - dataclasses are frozen in upstream implementation, so we reconstruct by type.
   546→
   547→        Args:
   548→            result: ForcedAlignResult
   549→            offset_sec: Offset in seconds
   550→
   551→        Returns:
   552→            ForcedAlignResult: New object with shifted timestamps.
   553→        """
   554→        if result is None:
   555→            return None
   556→        items = []
   557→        for it in result.items:
   558→            items.append(type(it)(text=it.text,
   559→                                  start_time=round(it.start_time + offset_sec, 3),
   560→                                  end_time=round(it.end_time + offset_sec, 3)))
   561→        return type(result)(items=items)
   562→
   563→    def _merge_align_results(self, results: List[Any]) -> Optional[Any]:
   564→        """
   565→        Merge multiple ForcedAlignResult objects into a single one by concatenating items.
   566→
   567→        Args:
   568→            results: List of ForcedAlignResult
   569→
   570→        Returns:
   571→            ForcedAlignResult or None
   572→        """
   573→        if not results:
   574→            return None
   575→        all_items = []
   576→        for r in results:
   577→            if r is None:
   578→                continue
   579→            all_items.extend(list(r.items))
   580→        if not all_items:
   581→            return None
   582→        return type(results[0])(items=all_items)
   583→
   584→    def init_streaming_state(
   585→        self,
   586→        context: str = "",
   587→        language: Optional[str] = None,
   588→        unfixed_chunk_num: int = 2,
   589→        unfixed_token_num: int = 5,
   590→        chunk_size_sec: float = 2.0,
   591→    ) -> ASRStreamingState:
   592→        """
   593→        Initialize streaming ASR state for a single stream.
   594→
   595→        Notes:
   596→            - Streaming ASR is supported ONLY for vLLM backend.
   597→            - Streaming ASR does NOT support timestamps (forced aligner is not used).
   598→            - Batch inference is NOT supported.
   599→
   600→        Args:
   601→            context:
   602→                Context string.
   603→            language:
   604→                Optional forced language. If provided, it must be in supported languages.
   605→                Same behavior as transcribe(): forces text-only output via prompt suffix.
   606→            unfixed_chunk_num:
   607→                For the first N chunks, do not use previous output as prefix prompt (reset prefix to "").
   608→            unfixed_token_num:
   609→                Roll back the last K tokens from accumulated output when using it as prefix prompt
   610→                after unfixed_chunk_num.
   611→            chunk_size_sec:
   612→                Chunk size in seconds (audio is 16k PCM). The function will internally convert it
   613→                to sample count at 16kHz.
   614→
   615→        Returns:
   616→            ASRStreamingState: Mutable state object to be passed to streaming_transcribe() and
   617→            finish_streaming_transcribe().
   618→
   619→        Raises:
   620→            ValueError:
   621→                - If backend is not "vllm".
   622→                - If chunk_size_sec <= 0.
   623→                - If forced language is invalid (same validation rules as transcribe()).
   624→        """
   625→        if self.backend != "vllm":
   626→            raise ValueError("Streaming ASR is supported only for vLLM backend (backend='vllm').")
   627→        if chunk_size_sec is None or float(chunk_size_sec) <= 0:
   628→            raise ValueError(f"chunk_size_sec must be > 0, got: {chunk_size_sec}")
   629→
   630→        force_language = None
   631→        if language is not None and str(language).strip() != "":
   632→            ln = normalize_language_name(str(language))
   633→            validate_language(ln)
   634→            force_language = ln
   635→
   636→        chunk_size_samples = int(round(float(chunk_size_sec) * SAMPLE_RATE))
   637→        chunk_size_samples = max(1, chunk_size_samples)
   638→
   639→        prompt_raw = self._build_text_prompt(context=context, force_language=force_language)
   640→
   641→        return ASRStreamingState(
   642→            unfixed_chunk_num=int(unfixed_chunk_num),
   643→            unfixed_token_num=int(unfixed_token_num),
   644→            chunk_size_sec=float(chunk_size_sec),
   645→            chunk_size_samples=int(chunk_size_samples),
   646→            chunk_id=0,
   647→            buffer=np.zeros((0,), dtype=np.float32),
   648→            audio_accum=np.zeros((0,), dtype=np.float32),
   649→            prompt_raw=prompt_raw,
   650→            context=context or "",
   651→            force_language=force_language,
   652→            language="",
   653→            text="",
   654→            _raw_decoded="",
   655→        )
   656→
   657→    def streaming_transcribe(self, pcm16k: np.ndarray, state: ASRStreamingState) -> ASRStreamingState:
   658→        """
   659→        Streaming ASR decode step.
   660→
   661→        This function accepts an arbitrary-length 16k PCM float numpy array (mono).
   662→        It buffers incoming samples, and whenever enough samples are accumulated to form one
   663→        full chunk (chunk_size_sec), it runs one incremental decode step and updates:
   664→
   665→          - state.language
   666→          - state.text
   667→
   668→        The caller only needs to keep passing audio to this function and read state.language/state.text.
   669→
   670→        Implementation details:
   671→            - Each time a new chunk is ready, we append it to audio_accum and re-feed *all* audio seen
   672→              so far to the model (no padding).
   673→            - We update the prompt as: state.prompt_raw + prefix_text
   674→            - Prefix rollback strategy:
   675→                * If chunk_id < unfixed_chunk_num: prefix_text = ""
   676→                * Else: rollback last unfixed_token_num tokens from previously accumulated decoded text.
   677→
   678→        Notes:
   679→            - vLLM backend only.
   680→            - No timestamps.
   681→            - Single stream only (no batching).
   682→
   683→        Args:
   684→            pcm16k:
   685→                16kHz mono PCM waveform (np.ndarray). Length can be any non-negative integer.
   686→                dtype can be float32/float64/int16; it will be converted to float32.
   687→            state:
   688→                Streaming state returned by init_streaming_state().
   689→
   690→        Returns:
   691→            ASRStreamingState: The same state object (mutated) for convenience.
   692→
   693→        Raises:
   694→            ValueError:
   695→                If backend is not "vllm" or state is invalid.
   696→        """
   697→        if self.backend != "vllm":
   698→            raise ValueError("streaming_transcribe() is supported only for vLLM backend (backend='vllm').")
   699→        if state is None:
   700→            raise ValueError("state must not be None. Call init_streaming_state() first.")
   701→        if pcm16k is None:
   702→            raise ValueError("pcm16k must not be None.")
   703→
   704→        # Ensure 1D mono
   705→        x = np.asarray(pcm16k)
   706→        if x.ndim != 1:
   707→            x = x.reshape(-1)
   708→
   709→        # Convert to float32 PCM in [-1, 1] if int16 provided
   710→        if x.dtype == np.int16:
   711→            x = (x.astype(np.float32) / 32768.0)
   712→        else:
   713→            x = x.astype(np.float32, copy=False)
   714→
   715→        # Append to buffer
   716→        if x.shape[0] > 0:
   717→            state.buffer = np.concatenate([state.buffer, x], axis=0)
   718→
   719→        # Consume full chunks
   720→        while state.buffer.shape[0] >= state.chunk_size_samples:
   721→            chunk = state.buffer[: state.chunk_size_samples]
   722→            state.buffer = state.buffer[state.chunk_size_samples :]
   723→
   724→            # Accumulate audio (re-feed from start, no padding)
   725→            if state.audio_accum.shape[0] == 0:
   726→                state.audio_accum = chunk
   727→            else:
   728→                state.audio_accum = np.concatenate([state.audio_accum, chunk], axis=0)
   729→
   730→            # Build prefix with rollback strategy
   731→            prefix = ""
   732→            if state.chunk_id < state.unfixed_chunk_num:
   733→                prefix = ""
   734→            else:
   735→                cur_ids = self.processor.tokenizer.encode(state._raw_decoded)
   736→                k = int(state.unfixed_token_num)
   737→                while True:
   738→                    end_idx = max(0, len(cur_ids) - k)
   739→                    prefix = self.processor.tokenizer.decode(cur_ids[:end_idx]) if end_idx > 0 else ""
   740→                    if '\ufffd' not in prefix:
   741→                        break
   742→                    else:
   743→                        if end_idx == 0:
   744→                            prefix = ""
   745→                            break
   746→                        k += 1
   747→
   748→            prompt = state.prompt_raw + prefix
   749→
   750→            # vLLM input: single item
   751→            inp = {"prompt": prompt, "multi_modal_data": {"audio": [state.audio_accum]}}
   752→
   753→            outputs = self.model.generate([inp], sampling_params=self.sampling_params, use_tqdm=False)
   754→            gen_text = outputs[0].outputs[0].text
   755→
   756→            # Accumulate raw decoded (then parse to lang/text)
   757→            state._raw_decoded = (prefix + gen_text) if prefix is not None else gen_text
   758→
   759→            lang, txt = parse_asr_output(state._raw_decoded, user_language=state.force_language)
   760→            state.language = lang
   761→            state.text = txt
   762→
   763→            state.chunk_id += 1
   764→
   765→        return state
   766→
   767→    def finish_streaming_transcribe(self, state: ASRStreamingState) -> ASRStreamingState:
   768→        """
   769→        Finish streaming ASR.
   770→
   771→        This function flushes the remaining buffered audio in state.buffer (tail audio).
   772→        It sends the remaining samples to the model even if shorter than chunk_size_sec,
   773→        without padding. Then it updates state.language/state.text one last time.
   774→
   775→        Notes:
   776→            - vLLM backend only.
   777→            - No timestamps.
   778→            - Single stream only.
   779→
   780→        Args:
   781→            state:
   782→                Streaming state.
   783→
   784→        Returns:
   785→            ASRStreamingState: Updated state (mutated).
   786→
   787→        Raises:
   788→            ValueError:
   789→                If backend is not "vllm" or state is invalid.
   790→        """
   791→        if self.backend != "vllm":
   792→            raise ValueError("finish_streaming_transcribe() is supported only for vLLM backend (backend='vllm').")
   793→        if state is None:
   794→            raise ValueError("state must not be None.")
   795→
   796→        # If no remaining buffer, still return state as-is.
   797→        if state.buffer is None or state.buffer.shape[0] == 0:
   798→            return state
   799→
   800→        tail = state.buffer
   801→        state.buffer = np.zeros((0,), dtype=np.float32)
   802→
   803→        # Append tail to accumulated audio
   804→        if state.audio_accum.shape[0] == 0:
   805→            state.audio_accum = tail
   806→        else:
   807→            state.audio_accum = np.concatenate([state.audio_accum, tail], axis=0)
   808→
   809→        # Prefix rollback strategy (same as per-chunk)
   810→        prefix = ""
   811→        if state.chunk_id < state.unfixed_chunk_num:
   812→            prefix = ""
   813→        else:
   814→            cur_ids = self.processor.tokenizer.encode(state._raw_decoded)
   815→            end_idx = max(1, len(cur_ids) - int(state.unfixed_token_num))
   816→            prefix = self.processor.tokenizer.decode(cur_ids[:end_idx])
   817→
   818→        prompt = state.prompt_raw + prefix
   819→        inp = {"prompt": prompt, "multi_modal_data": {"audio": [state.audio_accum]}}
   820→
   821→        outputs = self.model.generate([inp], sampling_params=self.sampling_params, use_tqdm=False)
   822→        gen_text = outputs[0].outputs[0].text
   823→
   824→        state._raw_decoded = (prefix + gen_text) if prefix is not None else gen_text
   825→        lang, txt = parse_asr_output(state._raw_decoded, user_language=state.force_language)
   826→        state.language = lang
   827→        state.text = txt
   828→
   829→        state.chunk_id += 1
   830→        return state

# coding=utf-8
# Copyright 2026 The Alibaba Qwen team.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch
from qwen_asr.core.transformers_backend import (
    Qwen3ASRConfig,
    Qwen3ASRForConditionalGeneration,
    Qwen3ASRProcessor,
)
from transformers import AutoConfig, AutoModel, AutoProcessor

AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)

from .qwen3_forced_aligner import Qwen3ForcedAligner
from .utils import (
    MAX_ASR_INPUT_SECONDS,
    MAX_FORCE_ALIGN_INPUT_SECONDS,
    SAMPLE_RATE,
    SUPPORTED_LANGUAGES,
    AudioChunk,
    AudioLike,
    chunk_list,
    merge_languages,
    normalize_audios,
    normalize_language_name,
    parse_asr_output,
    split_audio_into_chunks,
    validate_language,
)

try:
    from qwen_asr.core.vllm_backend import Qwen3ASRForConditionalGeneration
    from vllm import ModelRegistry
    ModelRegistry.register_model("Qwen3ASRForConditionalGeneration", Qwen3ASRForConditionalGeneration)
except:
    pass

@dataclass
class ASRTranscription:
    """
    One transcription result.

    Attributes:
        language (str):
            Merged language string for the sample, e.g. "Chinese" or "Chinese,English".
            Empty string if unknown or silent audio.
        text (str):
            Transcribed text.
        time_stamps (Optional[Any]):
            Forced aligner output (ForcedAlignResult).
            Present only when return_time_stamps=True.
    """
    language: str
    text: str
    time_stamps: Optional[Any] = None

@dataclass
class ASRStreamingState:
    """
    Streaming ASR state for one audio stream (single utterance).

    Attributes:
        unfixed_chunk_num (int):
            For the first N chunks, do not use previous ASR result as prefix prompt (reset prefix to "").
        unfixed_token_num (int):
            When chunk_id >= unfixed_chunk_num, rollback the last K tokens from the accumulated text
            before using it as prefix prompt, to reduce boundary jitter.
        chunk_size_sec (float):
            Chunk size in seconds. Audio will be fed to the model in increments of this length.
        chunk_size_samples (int):
            Chunk size in samples at 16kHz (derived from chunk_size_sec).
        chunk_id (int):
            Current chunk index (0-based).
        buffer (np.ndarray):
            Buffered PCM samples that are not yet consumed into a full chunk.
        audio_accum (np.ndarray):
            Accumulated audio from the beginning of the stream up to current time (no padding).
        prompt_raw (str):
            Base prompt generated by chat template (with generation prompt), without appended prefix text.
        context (str):
            Context string.
        force_language (Optional[str]):
            If provided, force output to be text-only by appending "language X<asr_text>" in prompt_raw,
            consistent with non-streaming transcribe().
        language (str):
            Latest parsed language (updated after each chunk decode). Empty if unknown/silent.
        text (str):
            Latest parsed transcription text (updated after each chunk decode).
        _raw_decoded (str):
            Internal accumulated decoded raw text (before parse_asr_output normalization).
            Used for rollback/token trimming and as prefix for prompting.
    """
    unfixed_chunk_num: int
    unfixed_token_num: int
    chunk_size_sec: float
    chunk_size_samples: int

    chunk_id: int
    buffer: np.ndarray
    audio_accum: np.ndarray

    prompt_raw: str
    context: str
    force_language: Optional[str]

    language: str
    text: str
    _raw_decoded: str

class Qwen3ASRModel:
    """
    Unified inference wrapper for Qwen3-ASR with two backends:
      - Transformers backend
      - vLLM backend

    It optionally supports time stamp output via Qwen3-ForcedAligner.

    Notes:
      - Each request uses a context text and exactly one audio.
      - If language is provided, the prompt will force the output to be text-only by appending
        "language {Language}<asr_text>" to the assistant prompt.
    """

    def __init__(
        self,
        backend: str,
        model: Any,
        processor: Any,
        sampling_params: Optional[Any] = None,
        forced_aligner: Optional[Qwen3ForcedAligner] = None,
        max_inference_batch_size: int = -1,
        max_new_tokens: int = 512,
    ):
        self.backend = backend  # "transformers" | "vllm"
        self.model = model
        self.processor = processor
        self.sampling_params = sampling_params
        self.forced_aligner = forced_aligner
        self.max_inference_batch_size = int(max_inference_batch_size)
        self.max_new_tokens = max_new_tokens

        if backend == "transformers":
            self.device = getattr(model, "device", None)
            if self.device is None:
                try:
                    self.device = next(model.parameters()).device
                except StopIteration:
                    self.device = torch.device("cpu")
            self.dtype = getattr(model, "dtype", torch.float32)
        else:
            self.device = None
            self.dtype = None

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str,
        forced_aligner: Optional[str] = None,
        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
        max_inference_batch_size: int = 32,
        max_new_tokens: Optional[int] = 512,
        **kwargs,
    ) -> "Qwen3ASRModel":
        """
        Initialize using Transformers backend.

        Args:
            pretrained_model_name_or_path:
                HuggingFace repo id or local directory.
            forced_aligner:
                Optional forced aligner model path/repo id.
            forced_aligner_kwargs:
                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
            max_inference_batch_size:
                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
            max_new_tokens:
                Maximum number of tokens to generate.
            **kwargs:
                Forwarded to AutoModel.from_pretrained(...).

        Returns:
            Qwen3ASRModel
        """

        model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)

        processor = AutoProcessor.from_pretrained(pretrained_model_name_or_path, fix_mistral_regex=True)

        forced_aligner_model = None
        if forced_aligner is not None:
            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
                forced_aligner, **(forced_aligner_kwargs or {})
            )

        return cls(
            backend="transformers",
            model=model,
            processor=processor,
            sampling_params=None,
            forced_aligner=forced_aligner_model,
            max_inference_batch_size=max_inference_batch_size,
            max_new_tokens=max_new_tokens,
        )

    @classmethod
    def LLM(
        cls,
        model: str,
        forced_aligner: Optional[str] = None,
        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
        max_inference_batch_size: int = -1,
        max_new_tokens: Optional[int] = 4096,
        **kwargs,
    ) -> "Qwen3ASRModel":
        """
        Initialize using vLLM backend.

        Import is isolated to keep vLLM optional.

        Args:
            model:
                Model path/repo for vLLM.
            forced_aligner:
                Optional forced aligner model path/repo id.
            forced_aligner_kwargs:
                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
            max_inference_batch_size:
                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
            max_new_tokens:
                Maximum number of tokens to generate.
            **kwargs:
                Forwarded to vllm.LLM(...).

        Returns:
            Qwen3ASRModel

        Raises:
            ImportError: If vLLM is not installed.
        """
        try:
            from vllm import LLM as vLLM
            from vllm import SamplingParams
        except Exception as e:
            raise ImportError(
                "vLLM is not available. Install with: pip install qwen-asr[vllm]"
            ) from e

        llm = vLLM(model=model, **kwargs)

        processor = Qwen3ASRProcessor.from_pretrained(model, fix_mistral_regex=True)
        sampling_params = SamplingParams(**({"temperature": 0.0, "max_tokens": max_new_tokens}))

        forced_aligner_model = None
        if forced_aligner is not None:
            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
                forced_aligner, **(forced_aligner_kwargs or {})
            )

        return cls(
            backend="vllm",
            model=llm,
            processor=processor,
            sampling_params=sampling_params,
            forced_aligner=forced_aligner_model,
            max_inference_batch_size=max_inference_batch_size,
            max_new_tokens=None,
        )

    def get_supported_languages(self) -> List[str]:
        """
        Returns the supported language list.

        Returns:
            List[str]: Canonical language names.
        """
        return list(SUPPORTED_LANGUAGES)

    @torch.no_grad()
    def transcribe(
        self,
        audio: Union[AudioLike, List[AudioLike]],
        context: Union[str, List[str]] = "",
        language: Optional[Union[str, List[Optional[str]]]] = None,
        return_time_stamps: bool = False,
    ) -> List[ASRTranscription]:
        """
        Transcribe audio with optional context and optional forced alignment timestamps.

        Args:
            audio:
                Audio input(s). Supported:
                  - str: local path / URL / base64 data url
                  - (np.ndarray, sr)
                  - list of above
            context:
                Context string(s). If scalar, it will be broadcast to batch size.
            language:
                Optional language(s). If provided, it must be in supported languages.
                If scalar, it will be broadcast to batch size.
                If provided, the prompt will force output to be transcription text only.
            return_time_stamps:
                If True, timestamps are produced via forced aligner and merged across chunks.
                This requires forced_aligner initialized.

        Returns:
            List[ASRTranscription]: One result per input audio.

        Raises:
            ValueError:
                - If return_time_stamps=True but forced_aligner is not provided.
                - If language is unsupported.
                - If batch sizes mismatch for context/language.
        """
        if return_time_stamps and self.forced_aligner is None:
            raise ValueError("return_time_stamps=True requires `forced_aligner` to be provided at initialization.")

        wavs = normalize_audios(audio)
        n = len(wavs)

        ctxs = context if isinstance(context, list) else [context]
        if len(ctxs) == 1 and n > 1:
            ctxs = ctxs * n
        if len(ctxs) != n:
            raise ValueError(f"Batch size mismatch: audio={n}, context={len(ctxs)}")

        langs_in: List[Optional[str]]
        if language is None:
            langs_in = [None] * n
        else:
            langs_in = language if isinstance(language, list) else [language]
            if len(langs_in) == 1 and n > 1:
                langs_in = langs_in * n
            if len(langs_in) != n:
                raise ValueError(f"Batch size mismatch: audio={n}, language={len(langs_in)}")

        langs_norm: List[Optional[str]] = []
        for l in langs_in:
            if l is None or str(l).strip() == "":
                langs_norm.append(None)
            else:
                ln = normalize_language_name(str(l))
                validate_language(ln)
                langs_norm.append(ln)

        max_chunk_sec = MAX_FORCE_ALIGN_INPUT_SECONDS if return_time_stamps else MAX_ASR_INPUT_SECONDS

        # chunk audios and record mapping
        chunks: List[AudioChunk] = []
        for i, wav in enumerate(wavs):
            parts = split_audio_into_chunks(
                wav=wav,
                sr=SAMPLE_RATE,
                max_chunk_sec=max_chunk_sec,
            )
            for j, (cwav, offset_sec) in enumerate(parts):
                chunks.append(AudioChunk(orig_index=i, chunk_index=j, wav=cwav, sr=SAMPLE_RATE, offset_sec=offset_sec))

        # run ASR on chunks
        chunk_ctx: List[str] = [ctxs[c.orig_index] for c in chunks]
        chunk_lang: List[Optional[str]] = [langs_norm[c.orig_index] for c in chunks]
        chunk_wavs: List[np.ndarray] = [c.wav for c in chunks]
        raw_outputs = self._infer_asr(chunk_ctx, chunk_wavs, chunk_lang)

        # parse outputs, prepare for optional alignment
        per_chunk_lang: List[str] = []
        per_chunk_text: List[str] = []
        for out, forced_lang in zip(raw_outputs, chunk_lang):
            lang, txt = parse_asr_output(out, user_language=forced_lang)
            per_chunk_lang.append(lang)
            per_chunk_text.append(txt)

        # forced alignment (optional)
        per_chunk_align: List[Optional[Any]] = [None] * len(chunks)
        if return_time_stamps:
            to_align_audio = []
            to_align_text = []
            to_align_lang = []
            to_align_idx = []

            for idx, (c, txt, lang_pred) in enumerate(zip(chunks, per_chunk_text, per_chunk_lang)):
                if txt.strip() == "":
                    continue
                to_align_audio.append((c.wav, c.sr))
                to_align_text.append(txt)
                to_align_lang.append(lang_pred)
                to_align_idx.append(idx)

            # batch align with max_inference_batch_size
            aligned_results: List[Any] = []
            for a_chunk, t_chunk, l_chunk in zip(
                chunk_list(to_align_audio, self.max_inference_batch_size),
                chunk_list(to_align_text, self.max_inference_batch_size),
                chunk_list(to_align_lang, self.max_inference_batch_size),
            ):
                aligned_results.extend(
                    self.forced_aligner.align(audio=a_chunk, text=t_chunk, language=l_chunk)
                )

            # offset fix
            for k, idx in enumerate(to_align_idx):
                c = chunks[idx]
                r = aligned_results[k]
                per_chunk_align[idx] = self._offset_align_result(r, c.offset_sec)

        # merge chunks back to original samples
        out_langs: List[List[str]] = [[] for _ in range(n)]
        out_texts: List[List[str]] = [[] for _ in range(n)]
        out_aligns: List[List[Any]] = [[] for _ in range(n)]

        for c, lang, txt, al in zip(chunks, per_chunk_lang, per_chunk_text, per_chunk_align):
            out_langs[c.orig_index].append(lang)
            out_texts[c.orig_index].append(txt)
            if return_time_stamps and al is not None:
                out_aligns[c.orig_index].append(al)

        results: List[ASRTranscription] = []
        for i in range(n):
            merged_text = "".join([t for t in out_texts[i] if t is not None])
            merged_language = merge_languages(out_langs[i])
            merged_align = None
            if return_time_stamps:
                merged_align = self._merge_align_results(out_aligns[i])
            results.append(ASRTranscription(language=merged_language, text=merged_text, time_stamps=merged_align))

        return results

    def _build_messages(self, context: str, audio_payload: Any) -> List[Dict[str, Any]]:
        return [
            {"role": "system", "content": context or ""},
            {"role": "user", "content": [{"type": "audio", "audio": audio_payload}]},
        ]

    def _build_text_prompt(self, context: str, force_language: Optional[str]) -> str:
        """
        Build the string prompt for one request.

        If force_language is provided, "language X<asr_text>" is appended after the generation prompt
        to request text-only output.
        """
        msgs = self._build_messages(context=context, audio_payload="")
        base = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
        if force_language:
            base = base + f"language {force_language}{'<asr_text>'}"
        return base

    def _infer_asr(
        self,
        contexts: List[str],
        wavs: List[np.ndarray],
        languages: List[Optional[str]],
    ) -> List[str]:
        """
        Run backend inference for chunk-level items.

        Args:
            contexts: List of context strings.
            wavs: List of mono waveforms (np.ndarray).
            languages: List of forced languages or None.

        Returns:
            List[str]: Raw decoded strings (one per chunk).
        """
        if self.backend == "transformers":
            return self._infer_asr_transformers(contexts, wavs, languages)
        if self.backend == "vllm":
            return self._infer_asr_vllm(contexts, wavs, languages)
        raise RuntimeError(f"Unknown backend: {self.backend}")

    def _infer_asr_transformers(
        self,
        contexts: List[str],
        wavs: List[np.ndarray],
        languages: List[Optional[str]],
    ) -> List[str]:
        outs: List[str] = []

        texts = [self._build_text_prompt(context=c, force_language=fl) for c, fl in zip(contexts, languages)]

        batch_size = self.max_inference_batch_size
        if batch_size is None or batch_size < 0:
            batch_size = len(texts)

        for i in range(0, len(texts), batch_size):
            sub_text = texts[i : i + batch_size]
            sub_wavs = wavs[i : i + batch_size]
            inputs = self.processor(text=sub_text, audio=sub_wavs, return_tensors="pt", padding=True)
            inputs = inputs.to(self.model.device).to(self.model.dtype)

            text_ids = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens)

            decoded = self.processor.batch_decode(
                text_ids.sequences[:, inputs["input_ids"].shape[1]:],
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            )
            outs.extend(list(decoded))

        return outs

    def _infer_asr_vllm(
        self,
        contexts: List[str],
        wavs: List[np.ndarray],
        languages: List[Optional[str]],
    ) -> List[str]:
        inputs: List[Dict[str, Any]] = []
        for c, w, fl in zip(contexts, wavs, languages):
            prompt = self._build_text_prompt(context=c, force_language=fl)
            inputs.append({"prompt": prompt, "multi_modal_data": {"audio": [w]}})

        outs: List[str] = []
        for batch in chunk_list(inputs, self.max_inference_batch_size):
            outputs = self.model.generate(batch, sampling_params=self.sampling_params, use_tqdm=False)
            for o in outputs:
                outs.append(o.outputs[0].text)
        return outs

    def _offset_align_result(self, result: Any, offset_sec: float) -> Any:
        """
        Apply time offset to a ForcedAlignResult-like object.

        This function assumes:
          - result has attribute `.items` which is a list of items with start_time/end_time in seconds.
          - dataclasses are frozen in upstream implementation, so we reconstruct by type.

        Args:
            result: ForcedAlignResult
            offset_sec: Offset in seconds

        Returns:
            ForcedAlignResult: New object with shifted timestamps.
        """
        if result is None:
            return None
        items = []
        for it in result.items:
            items.append(type(it)(text=it.text,
                                  start_time=round(it.start_time + offset_sec, 3),
                                  end_time=round(it.end_time + offset_sec, 3)))
        return type(result)(items=items)

    def _merge_align_results(self, results: List[Any]) -> Optional[Any]:
        """
        Merge multiple ForcedAlignResult objects into a single one by concatenating items.

        Args:
            results: List of ForcedAlignResult

        Returns:
            ForcedAlignResult or None
        """
        if not results:
            return None
        all_items = []
        for r in results:
            if r is None:
                continue
            all_items.extend(list(r.items))
        if not all_items:
            return None
        return type(results[0])(items=all_items)

    def init_streaming_state(
        self,
        context: str = "",
        language: Optional[str] = None,
        unfixed_chunk_num: int = 2,
        unfixed_token_num: int = 5,
        chunk_size_sec: float = 2.0,
    ) -> ASRStreamingState:
        """
        Initialize streaming ASR state for a single stream.

        Notes:
            - Streaming ASR is supported ONLY for vLLM backend.
            - Streaming ASR does NOT support timestamps (forced aligner is not used).
            - Batch inference is NOT supported.

        Args:
            context:
                Context string.
            language:
                Optional forced language. If provided, it must be in supported languages.
                Same behavior as transcribe(): forces text-only output via prompt suffix.
            unfixed_chunk_num:
                For the first N chunks, do not use previous output as prefix prompt (reset prefix to "").
            unfixed_token_num:
                Roll back the last K tokens from accumulated output when using it as prefix prompt
                after unfixed_chunk_num.
            chunk_size_sec:
                Chunk size in seconds (audio is 16k PCM). The function will internally convert it
                to sample count at 16kHz.

        Returns:
            ASRStreamingState: Mutable state object to be passed to streaming_transcribe() and
            finish_streaming_transcribe().

        Raises:
            ValueError:
                - If backend is not "vllm".
                - If chunk_size_sec <= 0.
                - If forced language is invalid (same validation rules as transcribe()).
        """
        if self.backend != "vllm":
            raise ValueError("Streaming ASR is supported only for vLLM backend (backend='vllm').")
        if chunk_size_sec is None or float(chunk_size_sec) <= 0:
            raise ValueError(f"chunk_size_sec must be > 0, got: {chunk_size_sec}")

        force_language = None
        if language is not None and str(language).strip() != "":
            ln = normalize_language_name(str(language))
            validate_language(ln)
            force_language = ln

        chunk_size_samples = int(round(float(chunk_size_sec) * SAMPLE_RATE))
        chunk_size_samples = max(1, chunk_size_samples)

        prompt_raw = self._build_text_prompt(context=context, force_language=force_language)

        return ASRStreamingState(
            unfixed_chunk_num=int(unfixed_chunk_num),
            unfixed_token_num=int(unfixed_token_num),
            chunk_size_sec=float(chunk_size_sec),
            chunk_size_samples=int(chunk_size_samples),
            chunk_id=0,
            buffer=np.zeros((0,), dtype=np.float32),
            audio_accum=np.zeros((0,), dtype=np.float32),
            prompt_raw=prompt_raw,
            context=context or "",
            force_language=force_language,
            language="",
            text="",
            _raw_decoded="",
        )

    def streaming_transcribe(self, pcm16k: np.ndarray, state: ASRStreamingState) -> ASRStreamingState:
        """
        Streaming ASR decode step.

        This function accepts an arbitrary-length 16k PCM float numpy array (mono).
        It buffers incoming samples, and whenever enough samples are accumulated to form one
        full chunk (chunk_size_sec), it runs one incremental decode step and updates:

          - state.language
          - state.text

        The caller only needs to keep passing audio to this function and read state.language/state.text.

        Implementation details:
            - Each time a new chunk is ready, we append it to audio_accum and re-feed *all* audio seen
              so far to the model (no padding).
            - We update the prompt as: state.prompt_raw + prefix_text
            - Prefix rollback strategy:
                * If chunk_id < unfixed_chunk_num: prefix_text = ""
                * Else: rollback last unfixed_token_num tokens from previously accumulated decoded text.

        Notes:
            - vLLM backend only.
            - No timestamps.
            - Single stream only (no batching).

        Args:
            pcm16k:
                16kHz mono PCM waveform (np.ndarray). Length can be any non-negative integer.
                dtype can be float32/float64/int16; it will be converted to float32.
            state:
                Streaming state returned by init_streaming_state().

        Returns:
            ASRStreamingState: The same state object (mutated) for convenience.

        Raises:
            ValueError:
                If backend is not "vllm" or state is invalid.
        """
        if self.backend != "vllm":
            raise ValueError("streaming_transcribe() is supported only for vLLM backend (backend='vllm').")
        if state is None:
            raise ValueError("state must not be None. Call init_streaming_state() first.")
        if pcm16k is None:
            raise ValueError("pcm16k must not be None.")

        # Ensure 1D mono
        x = np.asarray(pcm16k)
        if x.ndim != 1:
            x = x.reshape(-1)

        # Convert to float32 PCM in [-1, 1] if int16 provided
        if x.dtype == np.int16:
            x = (x.astype(np.float32) / 32768.0)
        else:
            x = x.astype(np.float32, copy=False)

        # Append to buffer
        if x.shape[0] > 0:
            state.buffer = np.concatenate([state.buffer, x], axis=0)

        # Consume full chunks
        while state.buffer.shape[0] >= state.chunk_size_samples:
            chunk = state.buffer[: state.chunk_size_samples]
            state.buffer = state.buffer[state.chunk_size_samples :]

            # Accumulate audio (re-feed from start, no padding)
            if state.audio_accum.shape[0] == 0:
                state.audio_accum = chunk
            else:
                state.audio_accum = np.concatenate([state.audio_accum, chunk], axis=0)

            # Build prefix with rollback strategy
            prefix = ""
            if state.chunk_id < state.unfixed_chunk_num:
                prefix = ""
            else:
                cur_ids = self.processor.tokenizer.encode(state._raw_decoded)
                k = int(state.unfixed_token_num)
                while True:
                    end_idx = max(0, len(cur_ids) - k)
                    prefix = self.processor.tokenizer.decode(cur_ids[:end_idx]) if end_idx > 0 else ""
                    if '\ufffd' not in prefix:
                        break
                    else:
                        if end_idx == 0:
                            prefix = ""
                            break
                        k += 1

            prompt = state.prompt_raw + prefix

            # vLLM input: single item
            inp = {"prompt": prompt, "multi_modal_data": {"audio": [state.audio_accum]}}

            outputs = self.model.generate([inp], sampling_params=self.sampling_params, use_tqdm=False)
            gen_text = outputs[0].outputs[0].text

            # Accumulate raw decoded (then parse to lang/text)
            state._raw_decoded = (prefix + gen_text) if prefix is not None else gen_text

            lang, txt = parse_asr_output(state._raw_decoded, user_language=state.force_language)
            state.language = lang
            state.text = txt

            state.chunk_id += 1

        return state

    def finish_streaming_transcribe(self, state: ASRStreamingState) -> ASRStreamingState:
        """
        Finish streaming ASR.

        This function flushes the remaining buffered audio in state.buffer (tail audio).
        It sends the remaining samples to the model even if shorter than chunk_size_sec,
        without padding. Then it updates state.language/state.text one last time.

        Notes:
            - vLLM backend only.
            - No timestamps.
            - Single stream only.

        Args:
            state:
                Streaming state.

        Returns:
            ASRStreamingState: Updated state (mutated).

        Raises:
            ValueError:
                If backend is not "vllm" or state is invalid.
        """
        if self.backend != "vllm":
            raise ValueError("finish_streaming_transcribe() is supported only for vLLM backend (backend='vllm').")
        if state is None:
            raise ValueError("state must not be None.")

        # If no remaining buffer, still return state as-is.
        if state.buffer is None or state.buffer.shape[0] == 0:
            return state

        tail = state.buffer
        state.buffer = np.zeros((0,), dtype=np.float32)

        # Append tail to accumulated audio
        if state.audio_accum.shape[0] == 0:
            state.audio_accum = tail
        else:
            state.audio_accum = np.concatenate([state.audio_accum, tail], axis=0)

        # Prefix rollback strategy (same as per-chunk)
        prefix = ""
        if state.chunk_id < state.unfixed_chunk_num:
            prefix = ""
        else:
            cur_ids = self.processor.tokenizer.encode(state._raw_decoded)
            end_idx = max(1, len(cur_ids) - int(state.unfixed_token_num))
            prefix = self.processor.tokenizer.decode(cur_ids[:end_idx])

        prompt = state.prompt_raw + prefix
        inp = {"prompt": prompt, "multi_modal_data": {"audio": [state.audio_accum]}}

        outputs = self.model.generate([inp], sampling_params=self.sampling_params, use_tqdm=False)
        gen_text = outputs[0].outputs[0].text

        state._raw_decoded = (prefix + gen_text) if prefix is not None else gen_text
        lang, txt = parse_asr_output(state._raw_decoded, user_language=state.force_language)
        state.language = lang
        state.text = txt

        state.chunk_id += 1
        return state


Bash

Check ckpt-24000 files


added_tokens.json
chat_template.json
config.json
generation_config.json
merges.txt
model.safetensors
preprocessor_config.json
special_tokens_map.json
tokenizer_config.json
tokenizer.json
vocab.json


Now I have full understanding. Let me create the benchmark script. First, check dependencies.


pip list 2>/dev/null | grep -i jiwer; pip list 2>/dev/null | grep -i indic; pip list 2>/dev/null | grep -i normalizer

Check if jiwer and indic-nlp are installed


Install jiwer for WER/CER computation


nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv 2>&1

Check GPU availability


jiwer                                    3.1.0
indic-nlp-library                        0.92
indic_numtowords                         1.1.0
indic-unified-parser                     1.0.6
indictrans                               1.2.3
bnunicodenormalizer                      0.1.7
charset-normalizer                       3.4.6
whisper_normalizer                       0.1.12


Requirement already satisfied: rapidfuzz>=3.9.7 in /home/ubuntu/.local/lib/python3.10/site-packages (from jiwer) (3.14.3)
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)


name, memory.total [MiB], memory.free [MiB]
NVIDIA A100-SXM4-80GB, 81920 MiB, 11037 MiB


Only 11GB free on one GPU. Let me check all GPUs.


nvidia-smi --query-gpu=index,name,memory.total,memory.used,memory.free --format=csv 2>&1

Check all GPUs


index, name, memory.total [MiB], memory.used [MiB], memory.free [MiB]
0, NVIDIA A100-SXM4-80GB, 81920 MiB, 70119 MiB, 11037 MiB


Only one GPU with 11GB free. The qwen3 model is ~3.8GB in bf16, so it should fit. Let me now create the benchmark script.


ToolSearch

Question: select:TaskCreate


Question: select:TaskCreate


Question: select:TaskUpdate


Question: select:TaskUpdate


TaskCreate

Create benchmark script for qwen3-asr checkpoints

Write a Python script that loads a qwen3-asr checkpoint, runs inference on the indic-asr-benchmark-6k dataset, computes WER/CER/space-normalized WER per language, saves results as JSON/CSV


Task #1 created successfully: Create benchmark script for qwen3-asr checkpoints


Run benchmark on ckpt-24000

Run the benchmark script on checkpoints/qwen3-asr-ckpt-24000 and collect results


Task #2 created successfully: Run benchmark on ckpt-24000

Run benchmark on ckpt-24000


Run benchmark on ckpt-72000

Run the benchmark script on checkpoints/qwen3-asr-ckpt-72000 and collect results


Task #3 created successfully: Run benchmark on ckpt-72000

Run benchmark on ckpt-72000


Analyze results and sample quality

Compare WER/CER across checkpoints and languages, inspect sample predictions for quality, determine if model is producing meaningful and improving results


Task #4 created successfully: Analyze results and sample quality

Analyze results and sample quality


Updated task #1 status


Now let me create the benchmark script. For ASR benchmarks, big labs typically compute:
- **WER** (Word Error Rate) — standard, word-level edit distance
- **CER** (Character Error Rate) — critical for Indic languages where word boundaries vary
- **Space-normalized WER** — yes, it's a real thing (collapse multiple spaces, strip punctuation)
- Some also compute **MER** (Match Error Rate) and look at insertion/deletion/substitution breakdowns


Write

#!/usr/bin/env python3
"""
Benchmark script for Qwen3-ASR checkpoints on indic-asr-benchmark-6k.

Usage:
    python benchmark_qwen3_asr.py \
        --checkpoint /path/to/checkpoint \
        --dataset BayAreaBoys/indic-asr-benchmark-6k \
        --output-dir ./benchmark_results/ckpt-24000 \
        --batch-size 8 \
        --device cuda:0

Metrics computed per language and overall:
    - WER  (Word Error Rate)
    - CER  (Character Error Rate)
    - WER-norm (space/punctuation-normalized WER)
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from datasets import load_dataset
from jiwer import cer, wer

# Add the qwen3 repo to path
QWEN3_REPO = "/home/ubuntu/training/qwen3-asr-1.7b-phase2-sft"
sys.path.insert(0, QWEN3_REPO)

from qwen_asr import Qwen3ASRModel

# ---------------------------------------------------------------------------
# Text normalization helpers
# ---------------------------------------------------------------------------

def normalize_text_for_wer(text: str) -> str:
    """Standard normalization: lowercase, collapse whitespace, strip punctuation."""
    text = text.strip()
    text = text.lower()
    # Remove punctuation (keep letters, numbers, spaces, and Indic chars)
    text = re.sub(r'[^\w\s]', '', text, flags=re.UNICODE)
    # Collapse whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def normalize_text_for_cer(text: str) -> str:
    """For CER: strip outer whitespace, collapse internal spaces."""
    text = text.strip()
    text = re.sub(r'\s+', ' ', text)
    return text

# ---------------------------------------------------------------------------
# Main benchmark logic
# ---------------------------------------------------------------------------

def run_benchmark(args):
    print(f"=" * 70)
    print(f"Qwen3-ASR Benchmark")
    print(f"Checkpoint: {args.checkpoint}")
    print(f"Dataset:    {args.dataset}")
    print(f"Device:     {args.device}")
    print(f"Batch size: {args.batch_size}")
    print(f"=" * 70)

    # Load dataset
    print("\n[1/4] Loading dataset...")
    hf_token = os.environ.get("HF_TOKEN", args.hf_token)
    ds = load_dataset(
        args.dataset,
        token=hf_token,
        cache_dir=args.cache_dir,
        split="train",
    )
    print(f"  Loaded {len(ds)} samples")

    # Load model
    print(f"\n[2/4] Loading model from {args.checkpoint}...")
    t0 = time.time()
    asr = Qwen3ASRModel.from_pretrained(
        args.checkpoint,
        dtype=torch.bfloat16,
        device_map=args.device,
        max_inference_batch_size=args.batch_size,
        max_new_tokens=args.max_new_tokens,
    )
    load_time = time.time() - t0
    print(f"  Model loaded in {load_time:.1f}s")

    # Run inference
    print(f"\n[3/4] Running inference on {len(ds)} samples...")
    os.makedirs(args.output_dir, exist_ok=True)

    all_results = []
    languages = sorted(set(ds["language"]))
    print(f"  Languages: {languages}")

    # Group samples by language for organized processing
    lang_indices = defaultdict(list)
    for i, lang in enumerate(ds["language"]):
        lang_indices[lang].append(i)

    total_processed = 0
    total_time = 0.0

    for lang in languages:
        indices = lang_indices[lang]
        lang_results = []
        t_lang_start = time.time()

        # Process in batches
        for batch_start in range(0, len(indices), args.batch_size):
            batch_indices = indices[batch_start:batch_start + args.batch_size]
            batch_samples = [ds[i] for i in batch_indices]

            # Extract audio arrays
            audio_inputs = []
            for s in batch_samples:
                audio = s["audio"]
                wav = np.array(audio["array"], dtype=np.float32)
                sr = audio["sampling_rate"]
                audio_inputs.append((wav, sr))

            # Run transcription
            t_batch = time.time()
            try:
                results = asr.transcribe(
                    audio=audio_inputs,
                    language=None,  # auto-detect
                    return_time_stamps=False,
                )
            except Exception as e:
                print(f"  ERROR in batch {batch_start} for {lang}: {e}")
                # Fall back to one-by-one
                results = []
                for ai in audio_inputs:
                    try:
                        r = asr.transcribe(audio=ai, language=None, return_time_stamps=False)
                        results.extend(r)
                    except Exception as e2:
                        print(f"    Single sample error: {e2}")
                        from qwen_asr.inference.qwen3_asr import ASRTranscription
                        results.append(ASRTranscription(language="", text=""))

            batch_time = time.time() - t_batch

            for j, (idx, result) in enumerate(zip(batch_indices, results)):
                sample = ds[idx]
                record = {
                    "id": sample["id"],
                    "language": sample["language"],
                    "lang_code": sample["lang_code"],
                    "source": sample["source"],
                    "duration": sample["duration"],
                    "reference": sample["reference"],
                    "hypothesis": result.text,
                    "detected_language": result.language,
                }
                lang_results.append(record)

            total_processed += len(batch_indices)
            if total_processed % 100 == 0 or total_processed == len(ds):
                elapsed = time.time() - t_lang_start
                print(f"  [{total_processed}/{len(ds)}] {lang}: {len(lang_results)}/{len(indices)} "
                      f"({elapsed:.1f}s)")

        lang_time = time.time() - t_lang_start
        total_time += lang_time
        all_results.extend(lang_results)

    print(f"\n  Total inference time: {total_time:.1f}s "
          f"({total_time/len(ds)*1000:.0f}ms/sample)")

    # Save raw predictions
    predictions_path = os.path.join(args.output_dir, "predictions.json")
    with open(predictions_path, "w", encoding="utf-8") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"  Predictions saved to {predictions_path}")

    # Compute metrics
    print(f"\n[4/4] Computing metrics...")
    metrics = compute_metrics(all_results)

    # Save metrics
    metrics_path = os.path.join(args.output_dir, "metrics.json")
    with open(metrics_path, "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)
    print(f"  Metrics saved to {metrics_path}")

    # Print summary table
    print_summary(metrics, args.checkpoint)

    # Save summary CSV
    save_csv(metrics, os.path.join(args.output_dir, "metrics_summary.csv"))

    return metrics, all_results

def compute_metrics(results):
    """Compute WER, CER, normalized WER per language and overall."""
    metrics = {}

    # Group by language
    by_lang = defaultdict(list)
    for r in results:
        by_lang[r["language"]].append(r)

    overall_refs_raw = []
    overall_hyps_raw = []
    overall_refs_norm = []
    overall_hyps_norm = []
    overall_refs_cer = []
    overall_hyps_cer = []

    for lang in sorted(by_lang.keys()):
        records = by_lang[lang]

        refs_raw = []
        hyps_raw = []
        refs_norm = []
        hyps_norm = []
        refs_cer = []
        hyps_cer = []

        empty_hyp_count = 0

        for r in records:
            ref = r["reference"]
            hyp = r["hypothesis"]

            if not hyp or hyp.strip() == "":
                empty_hyp_count += 1

            # For WER (raw)
            ref_r = ref.strip() if ref.strip() else " "
            hyp_r = hyp.strip() if hyp.strip() else " "
            refs_raw.append(ref_r)
            hyps_raw.append(hyp_r)

            # For normalized WER
            ref_n = normalize_text_for_wer(ref)
            hyp_n = normalize_text_for_wer(hyp)
            ref_n = ref_n if ref_n else " "
            hyp_n = hyp_n if hyp_n else " "
            refs_norm.append(ref_n)
            hyps_norm.append(hyp_n)

            # For CER
            ref_c = normalize_text_for_cer(ref)
            hyp_c = normalize_text_for_cer(hyp)
            ref_c = ref_c if ref_c else " "
            hyp_c = hyp_c if hyp_c else " "
            refs_cer.append(ref_c)
            hyps_cer.append(hyp_c)

        wer_raw = wer(refs_raw, hyps_raw)
        wer_norm = wer(refs_norm, hyps_norm)
        cer_val = cer(refs_cer, hyps_cer)

        metrics[lang] = {
            "n_samples": len(records),
            "wer": round(wer_raw * 100, 2),
            "wer_normalized": round(wer_norm * 100, 2),
            "cer": round(cer_val * 100, 2),
            "empty_hypotheses": empty_hyp_count,
        }

        overall_refs_raw.extend(refs_raw)
        overall_hyps_raw.extend(hyps_raw)
        overall_refs_norm.extend(refs_norm)
        overall_hyps_norm.extend(hyps_norm)
        overall_refs_cer.extend(refs_cer)
        overall_hyps_cer.extend(hyps_cer)

    # Overall
    metrics["OVERALL"] = {
        "n_samples": len(results),
        "wer": round(wer(overall_refs_raw, overall_hyps_raw) * 100, 2),
        "wer_normalized": round(wer(overall_refs_norm, overall_hyps_norm) * 100, 2),
        "cer": round(cer(overall_refs_cer, overall_hyps_cer) * 100, 2),
        "empty_hypotheses": sum(m["empty_hypotheses"] for m in metrics.values() if "empty_hypotheses" in m),
    }

    return metrics

def print_summary(metrics, checkpoint_name):
    """Print a formatted summary table."""
    ckpt_label = Path(checkpoint_name).name
    print(f"\n{'=' * 70}")
    print(f"  BENCHMARK RESULTS — {ckpt_label}")
    print(f"{'=' * 70}")
    print(f"{'Language':<15} {'N':>6} {'WER%':>8} {'WER-norm%':>10} {'CER%':>8} {'Empty':>6}")
    print(f"{'-' * 55}")

    for lang in sorted(k for k in metrics if k != "OVERALL"):
        m = metrics[lang]
        print(f"{lang:<15} {m['n_samples']:>6} {m['wer']:>8.2f} {m['wer_normalized']:>10.2f} "
              f"{m['cer']:>8.2f} {m['empty_hypotheses']:>6}")

    print(f"{'-' * 55}")
    m = metrics["OVERALL"]
    print(f"{'OVERALL':<15} {m['n_samples']:>6} {m['wer']:>8.2f} {m['wer_normalized']:>10.2f} "
          f"{m['cer']:>8.2f} {m['empty_hypotheses']:>6}")
    print(f"{'=' * 70}")

def save_csv(metrics, path):
    """Save metrics as CSV."""
    with open(path, "w") as f:
        f.write("language,n_samples,wer,wer_normalized,cer,empty_hypotheses\n")
        for lang in sorted(k for k in metrics if k != "OVERALL"):
            m = metrics[lang]
            f.write(f"{lang},{m['n_samples']},{m['wer']},{m['wer_normalized']},{m['cer']},{m['empty_hypotheses']}\n")
        m = metrics["OVERALL"]
        f.write(f"OVERALL,{m['n_samples']},{m['wer']},{m['wer_normalized']},{m['cer']},{m['empty_hypotheses']}\n")
    print(f"  CSV saved to {path}")

def main():
    parser = argparse.ArgumentParser(description="Benchmark Qwen3-ASR on indic-asr-benchmark-6k")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--output-dir", default=None,
                        help="Output directory for results (default: benchmark_results/<ckpt-name>)")
    parser.add_argument("--batch-size", type=int, default=8, help="Inference batch size")
    parser.add_argument("--max-new-tokens", type=int, default=256, help="Max new tokens for generation")
    parser.add_argument("--device", default="cuda:0", help="Device to use")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k",
                        help="Dataset cache directory")

    args = parser.parse_args()

    if args.output_dir is None:
        ckpt_name = Path(args.checkpoint).name
        args.output_dir = f"/home/ubuntu/training/benchmark_results/{ckpt_name}"

    run_benchmark(args)

if __name__ == "__main__":
    main()


<tool_use_error>File has not been read yet. Read it first before writing to it.</tool_use_error>

Error: File has not been read yet. Read it first before writing to it.


Bash

ls /home/ubuntu/training/benchmark_qwen3_asr.py 2>&1 || echo "NOT_EXISTS"

Check if file exists