   300→# Keep legacy TarAudioLRUCache as fallback when indices don't exist
   301→class TarAudioLRUCache:
   302→    def __init__(self, max_open_tars: int = 16):
   303→        self.max_open_tars = int(max(1, max_open_tars))
   304→        self._cache: "OrderedDict[str, Any]" = OrderedDict()
   305→
   306→    def close(self):
   307→        import tarfile as _tarfile
   308→        for tf in self._cache.values():
   309→            try:
   310→                tf.close()
   311→            except Exception:
   312→                pass
   313→        self._cache.clear()
   314→
   315→    def _candidate_member_names(self, name: str) -> List[str]:
   316→        base = name.lstrip("./")
   317→        out = [name]
   318→        if not name.startswith("./"):
   319→            out.append(f"./{name}")
   320→        if base != name:
   321→            out.append(base)
   322→        return out
   323→
   324→    def _get_tar(self, tar_path: str):
   325→        import tarfile as _tarfile
   326→        real_path = os.path.realpath(tar_path)
   327→        tf = self._cache.get(real_path)
   328→        if tf is not None:
   329→            self._cache.move_to_end(real_path)
   330→            return tf
   331→        tf = _tarfile.open(real_path, mode="r:")
   332→        self._cache[real_path] = tf
   333→        if len(self._cache) > self.max_open_tars:
   334→            _, old_tf = self._cache.popitem(last=False)
   335→            try:
   336→                old_tf.close()
   337→            except Exception:
   338→                pass
   339→        return tf
   340→
   341→    def read_member(self, tar_path: str, member_name: str) -> bytes:
   342→        tf = self._get_tar(tar_path)
   343→        for cand in self._candidate_member_names(member_name):
   344→            try:
   345→                fobj = tf.extractfile(cand)
   346→            except KeyError:
   347→                fobj = None
   348→            if fobj is not None:
   349→                return fobj.read()
   350→        raise FileNotFoundError(f"Member not found: {member_name} in {tar_path}")
   351→
   352→
   353→class HybridTarReader:
   354→    """Uses IndexedTarReader when index exists, falls back to TarAudioLRUCache."""
   355→
   356→    def __init__(self, max_open_files: int = 64):
   357→        self._indexed = IndexedTarReader(max_open_files=max_open_files)
   358→        self._legacy = TarAudioLRUCache(max_open_tars=max(1, max_open_files // 4))
   359→        self._has_index: Dict[str, bool] = {}
   360→
   361→    def _check_index(self, tar_path: str) -> bool:
   362→        real_path = os.path.realpath(tar_path)
   363→        if real_path not in self._has_index:
   364→            self._has_index[real_path] = os.path.exists(real_path + ".index.json")
   365→        return self._has_index[real_path]
   366→
   367→    def read_member(self, tar_path: str, member_name: str) -> bytes:
   368→        if self._check_index(tar_path):
   369→            return self._indexed.read_member(tar_path, member_name)
   370→        return self._legacy.read_member(tar_path, member_name)
   371→
   372→    def close(self):
   373→        self._indexed.close()
   374→        self._legacy.close()
   375→
   376→
   377→def _make_tar_reader(max_open: int = 64, use_indexed: bool = True) -> Any:
   378→    """Create the best available tar reader."""
   379→    if use_indexed:
   380→        return HybridTarReader(max_open_files=max_open)
   381→    return TarAudioLRUCache(max_open_tars=max_open)
   382→
   383→
   384→def _decode_audio_bytes(raw: bytes, target_sr: int = 16000) -> Optional[np.ndarray]:
   385→    """Decode raw audio bytes to float32 numpy array at target_sr."""
   386→    try:
   387→        wav, sr = sf.read(io.BytesIO(raw), dtype="float32", always_2d=False)
   388→        wav = np.asarray(wav, dtype=np.float32)
   389→        if wav.ndim == 2:
   390→            wav = wav.mean(axis=1).astype(np.float32)
   391→        if sr != target_sr:
   392→            wav_t = torch.from_numpy(wav).unsqueeze(0)
   393→            wav_t = torchaudio.functional.resample(wav_t, sr, target_sr)
   394→            wav = wav_t.squeeze(0).numpy()
   395→        return wav
   396→    except Exception:
   397→        return None
   398→
   399→
   400→# ---------------------------------------------------------------------------
   401→# BucketedIterableDataset: reads per-bucket parquets, yields decoded waveforms
   402→# Worker-side decode: each worker decodes audio, collator only pads+tokenizes
   403→# ---------------------------------------------------------------------------
   404→
   405→class BucketedIterableDataset(IterableDataset):
   406→    """Reads from per-bucket parquet files, decodes audio in workers.
   407→
   408→    Yields PRE-FORMED BATCHES (lists of sample dicts) with per-bucket batch sizes.
   409→    The DataLoader should use batch_size=1 and the collator unwraps the inner list.
   410→
   411→    Key features:
   412→    - Same-duration batching (zero cross-bucket padding)
   413→    - Dynamic per-bucket batch sizes (short audio = big BS, long audio = small BS)
   414→    - Worker-side audio decoding (parallelized across num_workers)
   415→    - IndexedTarReader for O(1) file access
   416→    - Weighted bucket sampling proportional to audio hours
   417→    """
   418→
   419→    # Default per-bucket batch sizes (calibrated for H200 143GB, no GC)
   420→    # Memory: ~0.77 GB per (sample × audio_second). Target ≤125GB worst-case peak.
   421→    # OOM at b_10_15 BS=16 (152GB needed). Calibrated with 18GB safety margin.
   422→    DEFAULT_BUCKET_BS = {
   423→        "b_0_3":  56,   # 0-3s: worst 56×3×0.77=129GB
   424→        "b_3_5":  32,   # 3-5s: worst 32×5×0.77=123GB
   425→        "b_5_7":  24,   # 5-7s: worst 24×7×0.77=129GB
   426→        "b_7_10": 16,   # 7-10s: worst 16×10×0.77=123GB (34% of hours)
   427→        "b_10_15": 10,  # 10-15s: worst 10×15×0.77=116GB (30% of hours)
   428→        "b_15_20": 8,   # 15-20s: worst 8×20×0.77=123GB
   429→        "b_20_30": 4,   # 20-30s: worst 4×30×0.77=92GB
   430→    }
   431→
   432→    def __init__(
   433→        self,
   434→        bucket_dir: str,
   435→        bucket_config: dict,
   436→        prompt: str = "",
   437→        seed: int = 42,
   438→        manifest_batch_rows: int = 65536,
   439→        shuffle_row_groups: bool = True,
   440→        shuffle_within_batch: bool = True,
   441→        skip_empty_transcript: bool = True,
   442→        max_open_tars: int = 64,
   443→        use_indexed_tar: bool = True,
   444→        target_sr: int = 16000,
   445→        bucket_bs: Optional[Dict[str, int]] = None,
   446→        default_batch_size: int = 16,
   447→    ):
   448→        self.bucket_dir = bucket_dir
   449→        self.prompt = prompt
   450→        self.seed = int(seed)
   451→        self.manifest_batch_rows = int(manifest_batch_rows)
   452→        self.shuffle_row_groups = bool(shuffle_row_groups)
   453→        self.shuffle_within_batch = bool(shuffle_within_batch)
   454→        self.skip_empty_transcript = bool(skip_empty_transcript)
   455→        self.max_open_tars = max_open_tars
   456→        self.use_indexed_tar = use_indexed_tar
   457→        self.target_sr = target_sr
   458→        self.default_batch_size = default_batch_size
   459→        self.epoch = 0
   460→
   461→        # Per-bucket batch sizes
   462→        _bucket_bs = dict(self.DEFAULT_BUCKET_BS)
   463→        if bucket_bs:
   464→            _bucket_bs.update(bucket_bs)
   465→
   466→        # Parse bucket config
   467→        self.buckets = []
   468→        total_samples = 0
   469→        for b in bucket_config["buckets"]:
   470→            path = os.path.join(bucket_dir, f"{b['bucket_id']}.parquet")
   471→            if not os.path.exists(path):
   472→                continue
   473→            bid = b["bucket_id"]
   474→            self.buckets.append({
   475→                "id": bid,
   476→                "path": path,
   477→                "weight": b.get("hours", b.get("pct_hours", 1.0)),
   478→                "samples": b["samples"],
   479→                # Support both {bid: int} and {bid: {batch_size: int}} formats
   480→                "batch_size": (lambda v: v["batch_size"] if isinstance(v, dict) else v)(
   481→                    _bucket_bs.get(bid, default_batch_size)
   482→                ),
   483→            })
   484→            total_samples += b["samples"]
   485→
   486→        if not self.buckets:
   487→            raise ValueError(
   488→                f"No bucket parquet files found in {bucket_dir}. "
   489→                "Expected files like b_0_3.parquet alongside bucket_config.json."
   490→            )
   491→
   492→        self._estimated_num_rows = total_samples
   493→
   494→        # Normalize weights
   495→        total_weight = sum(b["weight"] for b in self.buckets)
   496→        for b in self.buckets:
   497→            b["norm_weight"] = b["weight"] / total_weight if total_weight > 0 else 1.0 / len(self.buckets)
   498→
   499→    @property
   500→    def estimated_num_rows(self) -> int:
   501→        return self._estimated_num_rows
   502→
   503→    def set_epoch(self, epoch: int):
   504→        self.epoch = int(epoch)
   505→
   506→    def __iter__(self):
   507→        """Yields pre-formed batches with dynamic per-bucket batch sizes.
   508→
   509→        Each yielded item is a list of sample dicts from a single bucket.
   510→        DataLoader should use batch_size=1 and collator unwraps the inner list.
   511→        """
   512→        worker_info = get_worker_info()
   513→        worker_id = worker_info.id if worker_info is not None else 0
   514→        num_workers = worker_info.num_workers if worker_info is not None else 1
   515→        rank, world_size = get_dist_info()
   516→        shard_id = rank * num_workers + worker_id
   517→        num_shards = world_size * num_workers
   518→
   519→        # Shared seed across shards keeps bucket selection order deterministic.
   520→        # This avoids cross-rank step drift when per-bucket BS differs.
   521→        rng = random.Random(self.seed + self.epoch * 1000003)
   522→
   523→        tar_reader = _make_tar_reader(self.max_open_tars, self.use_indexed_tar)
   524→
   525→        columns = ["sample_id", "language", "transcript", "tar_path",
   526→                    "tar_member_name", "duration_s", "bucket_id"]
   527→
   528→        try:
   529→            bucket_iters = []
   530→            for b in self.buckets:
   531→                bucket_iters.append({
   532→                    "info": b,
   533→                    "iter": self._iter_bucket(
   534→                        b["path"], columns, shard_id, num_shards, tar_reader
   535→                    ),
   536→                    "exhausted": False,
   537→                    "buffer": [],
   538→                })
   539→
   540→            weights = [b["info"]["norm_weight"] for b in bucket_iters]
   541→
   542→            while True:
   543→                active = [(i, bi) for i, bi in enumerate(bucket_iters) if not bi["exhausted"]]
   544→                if not active:
   545→                    break
   546→
   547→                active_weights = [weights[i] for i, _ in active]
   548→                total_w = sum(active_weights)
   549→                if total_w <= 0:
   550→                    break
   551→
   552→                chosen_idx = rng.choices([i for i, _ in active], weights=active_weights, k=1)[0]
   553→                bi = bucket_iters[chosen_idx]
   554→                target_bs = int(max(1, bi["info"]["batch_size"]))
   555→
   556→                # Fill buffer for this bucket to target dynamic BS.
   557→                while len(bi["buffer"]) < target_bs:
   558→                    sample = next(bi["iter"], None)
   559→                    if sample is None:
   560→                        bi["exhausted"] = True
   561→                        break
   562→                    bi["buffer"].append(sample)
   563→
   564→                if bi["buffer"]:
   565→                    batch = bi["buffer"][:target_bs]
   566→                    bi["buffer"] = bi["buffer"][target_bs:]
   567→                    yield batch
   568→
   569→            # Flush partial tail batches.
   570→            for bi in bucket_iters:
   571→                if bi["buffer"]:
   572→                    yield bi["buffer"]
   573→                    bi["buffer"] = []
   574→
   575→        finally:
   576→            tar_reader.close()
   577→
   578→    def _iter_bucket(self, parquet_path, columns, shard_id, num_shards, tar_reader):
   579→        """Iterate a single bucket's parquet, yielding decoded samples.
   580→
   581→        Sharding is done at the SAMPLE level (not row-group level) because
   582→        bucket parquets may have fewer row groups than total shards (GPUs × workers).
   583→        Every worker reads all row groups but only processes every num_shards-th sample.
   584→
   585→        IMPORTANT: Row group order and within-batch shuffle must be IDENTICAL across
   586→        all shards so that global_sample_idx maps to the same physical sample everywhere.
   587→        We use a SHARED seed (not shard-specific) for these shuffles.
   588→        """
   589→        pf = pq.ParquetFile(parquet_path)
   590→        num_rg = pf.num_row_groups
   591→
   592→        row_groups = list(range(num_rg))
   593→        # Shared RNG across all shards for deterministic row group order
   594→        shared_rng = random.Random(self.seed + self.epoch * 1000003)
   595→        if self.shuffle_row_groups:
   596→            shared_rng.shuffle(row_groups)
   597→
   598→        global_sample_idx = 0
   599→        for rg_idx in row_groups:
   600→            batches = pf.iter_batches(
   601→                batch_size=self.manifest_batch_rows,
   602→                row_groups=[rg_idx],
   603→                columns=columns,
   604→                use_threads=True,
   605→            )
   606→            for batch in batches:
   607→                data = {name: batch.column(i).to_pylist()
   608→                        for i, name in enumerate(batch.schema.names)}
   609→                n = batch.num_rows
   610→                indices = list(range(n))
   611→                if self.shuffle_within_batch:
   612→                    shared_rng.shuffle(indices)  # shared across shards for consistent global_sample_idx
   613→
   614→                for i in indices:
   615→                    # Sample-level sharding: each worker processes every num_shards-th sample
   616→                    if global_sample_idx % num_shards != shard_id:
   617→                        global_sample_idx += 1
   618→                        continue
   619→                    global_sample_idx += 1
   620→
   621→                    transcript = _as_str(data["transcript"][i]).strip()
   622→                    if self.skip_empty_transcript and not transcript:
   623→                        continue
   624→
   625→                    tar_path = _as_str(data["tar_path"][i])
   626→                    tar_member_name = _as_str(data["tar_member_name"][i])
   627→                    if not tar_path or not tar_member_name:
   628→                        continue
   629→
   630→                    # Clean transcript
   631→                    transcript = clean_transcript(transcript)
   632→                    if not transcript:
   633→                        continue
   634→
   635→                    # Detect script language
   636→                    declared_lang = _as_str(data["language"][i])
   637→                    effective_lang = detect_script_language(transcript, declared_lang)
   638→
   639→                    # WORKER-SIDE AUDIO DECODE
   640→                    try:
   641→                        raw_bytes = tar_reader.read_member(tar_path, tar_member_name)
   642→                        wav = _decode_audio_bytes(raw_bytes, self.target_sr)
   643→                        if wav is None or len(wav) == 0:
   644→                            continue
   645→                    except Exception:
   646→                        continue
   647→
   648→                    duration_s = data["duration_s"][i]
   649→                    bucket_id = _as_str(data.get("bucket_id", [""])[i]) if "bucket_id" in data else ""
   650→
   651→                    yield {
   652→                        "sample_id": _as_str(data["sample_id"][i]),
   653→                        "waveform": wav,
   654→                        "transcript": transcript,
   655→                        "language": effective_lang,
   656→                        "prompt": self.prompt,
   657→                        "duration_s": float(duration_s) if duration_s is not None else len(wav) / self.target_sr,
   658→                        "bucket_id": bucket_id,
   659→                        "tar_path": tar_path,
   660→                        "tar_member_name": tar_member_name,
   661→                    }
   662→
   663→
   664→# Legacy dataset for non-bucketed mode (backward compatible)
   665→class ParquetTarIterableDataset(IterableDataset):
   666→    def __init__(
   667→        self,
   668→        parquet_path: str,
   669→        split: str = "",
   670→        prompt: str = "",
   671→        seed: int = 42,
   672→        manifest_batch_rows: int = 65536,
   673→        shuffle_row_groups: bool = True,
   674→        shuffle_within_batch: bool = True,
   675→        min_duration_s: float = 0.0,
   676→        max_duration_s: float = 1200.0,
   677→        skip_empty_transcript: bool = True,
   678→        max_samples: int = 0,
   679→        max_open_tars: int = 64,
   680→        use_indexed_tar: bool = True,
   681→        target_sr: int = 16000,
   682→        worker_decode: bool = True,
   683→    ):
   684→        self.parquet_path = parquet_path
   685→        self.split = split.strip()
   686→        self.prompt = prompt
   687→        self.seed = int(seed)
   688→        self.manifest_batch_rows = int(manifest_batch_rows)
   689→        self.shuffle_row_groups = bool(shuffle_row_groups)
   690→        self.shuffle_within_batch = bool(shuffle_within_batch)
   691→        self.min_duration_s = float(min_duration_s)
   692→        self.max_duration_s = float(max_duration_s)
   693→        self.skip_empty_transcript = bool(skip_empty_transcript)
   694→        self.max_samples = int(max_samples)
   695→        self.max_open_tars = max_open_tars
   696→        self.use_indexed_tar = use_indexed_tar
   697→        self.target_sr = target_sr
   698→        self.worker_decode = worker_decode
   699→        self.epoch = 0


Found 2 files
qwen3-asr-1.7b-phase2-sft/finetuning/qwen3_asr_sft_phase2.py
qwen3-asr-1.7b-phase2-sft/finetuning/qwen3_asr_sft.py


Bash

grep -n "class.*Collat" /home/ubuntu/training/qwen3-asr-1.7b-phase2-sft/finetuning/*.py


   861→class DataCollatorForQwen3ASRPhase2:
   862→    """Collator that handles both worker-decoded waveforms and legacy tar-path mode."""
   863→    processor: Any
   864→    sampling_rate: int = 16000
   865→    max_open_tars: int = 16
   866→    language_tag_mode: str = "auto"
   867→    transcript_has_prefix: bool = False
   868→    skip_decode_errors: bool = True
   869→    worker_decode: bool = True
   870→    max_batch_seq_len: int = 500
   871→
   872→    def __post_init__(self):
   873→        # Only create tar cache if collator needs to decode (legacy mode)
   874→        if not self.worker_decode:
   875→            self._tar_cache = TarAudioLRUCache(max_open_tars=self.max_open_tars)
   876→        self._prefix_cache: Dict[str, str] = {}
   877→
   878→    def _make_zero_loss_batch(
   879→        self,
   880→        batch_debug: Optional[Dict[str, Any]] = None,
   881→        skip_reason: str = "",
   882→    ) -> Dict[str, Any]:
   883→        """Return a dummy batch that contributes zero loss but preserves breadcrumbs."""
   884→        debug = dict(batch_debug or {})
   885→        dummy_text = self._build_prefix_text("") + "language None<asr_text>dummy" + (
   886→            self.processor.tokenizer.eos_token or ""
   887→        )
   888→        dummy_wav = np.zeros(16000, dtype=np.float32)
   889→        full_inputs = self.processor(
   890→            text=[dummy_text],
   891→            audio=[dummy_wav],
   892→            return_tensors="pt",
   893→            padding=True,
   894→            truncation=False,
   895→        )
   896→        labels = full_inputs["input_ids"].clone()
   897→        full_inputs["labels"] = labels
   898→        full_inputs["loss_scale"] = torch.tensor(0.0, dtype=torch.float32)
   899→        full_inputs["_batch_size"] = torch.tensor(int(debug.get("sample_count", 0)), dtype=torch.long)
   900→        full_inputs["_batch_audio_secs"] = torch.tensor(
   901→            float(debug.get("audio_secs", 0.0)), dtype=torch.float32
   902→        )
   903→        full_inputs["_batch_seq_len"] = torch.tensor(int(debug.get("seq_len", 0)), dtype=torch.long)
   904→        full_inputs["_batch_attn_tokens_sum"] = torch.tensor(
   905→            int(debug.get("attn_tokens_sum", 0)), dtype=torch.long
   906→        )
   907→        full_inputs["_batch_audio_tokens_max"] = torch.tensor(
   908→            int(debug.get("audio_tokens_max", 0)), dtype=torch.long
   909→        )
   910→        full_inputs["_batch_label_tokens_sum"] = torch.tensor(
   911→            int(debug.get("label_tokens_sum", 0)), dtype=torch.long
   912→        )
   913→        full_inputs["_batch_skip_reason"] = skip_reason
   914→        full_inputs["_batch_debug_json"] = json.dumps(
   915→            debug, ensure_ascii=True, separators=(",", ":")
   916→        )
   917→        return full_inputs
   918→
   919→    def _build_prefix_text(self, prompt: str) -> str:
   920→        if prompt in self._prefix_cache:
   921→            return self._prefix_cache[prompt]
   922→        prefix_msgs = build_prefix_messages(prompt, None)
   923→        prefix_text = self.processor.apply_chat_template(
   924→            [prefix_msgs], add_generation_prompt=True, tokenize=False
   925→        )[0]
   926→        self._prefix_cache[prompt] = prefix_text
   927→        return prefix_text
   928→
   929→    def _decode_audio_legacy(self, tar_path: str, member_name: str) -> np.ndarray:
   930→        """Legacy: decode audio from tar in collator (main process)."""
   931→        raw = self._tar_cache.read_member(tar_path, member_name)
   932→        wav, sr = sf.read(io.BytesIO(raw), dtype="float32", always_2d=False)
   933→        wav = np.asarray(wav, dtype=np.float32)
   934→        if wav.ndim == 2:
   935→            wav = wav.mean(axis=1).astype(np.float32)
   936→        if sr != self.sampling_rate:
   937→            wav_t = torch.from_numpy(wav).unsqueeze(0)
   938→            wav_t = torchaudio.functional.resample(wav_t, sr, self.sampling_rate)
   939→            wav = wav_t.squeeze(0).numpy()
   940→        return wav
   941→
   942→    def __call__(self, features: List[Any]) -> Dict[str, torch.Tensor]:
   943→        # Bucketed dataset emits pre-formed batches. DataLoader wraps that once
   944→        # when batch_size=1, so unwrap here.
   945→        if features and isinstance(features[0], list):
   946→            features = features[0]
   947→
   948→        audios: List[np.ndarray] = []
   949→        prefix_texts: List[str] = []
   950→        targets: List[str] = []
   951→        sample_debug: List[Dict[str, Any]] = []
   952→
   953→        for f in features:
   954→            if self.worker_decode and "waveform" in f:
   955→                # Worker already decoded the audio
   956→                wav = f["waveform"]
   957→                if wav is None or len(wav) == 0:
   958→                    continue
   959→                transcript = f.get("transcript", "").strip()
   960→                language = f.get("language", "")
   961→            else:
   962→                # Legacy: decode in collator
   963→                try:
   964→                    wav = self._decode_audio_legacy(f["tar_path"], f["tar_member_name"])
   965→                except Exception:
   966→                    if self.skip_decode_errors:
   967→                        continue
   968→                    raise
   969→                transcript = clean_transcript(_as_str(f.get("transcript", "")))
   970→                if not transcript:
   971→                    continue
   972→                declared_lang = _as_str(f.get("language", ""))
   973→                language = detect_script_language(transcript, declared_lang)
   974→
   975→            if not transcript:
   976→                continue
   977→
   978→            # Sanitize audio: skip samples with NaN/Inf or extreme length
   979→            wav_arr = np.asarray(wav)
   980→            if np.isnan(wav_arr).any() or np.isinf(wav_arr).any():
   981→                import warnings
   982→                sid = f.get("sample_id", "?")
   983→                warnings.warn(f"Skipping sample {sid}: audio contains NaN/Inf")
   984→                continue
   985→            if len(wav_arr) > 30 * self.sampling_rate:  # >30s at sample rate
   986→                import warnings
   987→                sid = f.get("sample_id", "?")
   988→                warnings.warn(f"Skipping sample {sid}: audio too long ({len(wav_arr)/self.sampling_rate:.1f}s)")
   989→                continue
   990→
   991→            target = format_target_text(
   992→                transcript=transcript,
   993→                language_code=language,
   994→                language_tag_mode=self.language_tag_mode,
   995→                transcript_has_prefix=self.transcript_has_prefix,
   996→            )
   997→
   998→            prompt = _as_str(f.get("prompt", ""))
   999→            prefix_text = self._build_prefix_text(prompt)
  1000→            duration_s = f.get("duration_s")
  1001→            if duration_s is None:
  1002→                duration_s = len(wav_arr) / self.sampling_rate
  1003→
  1004→            audios.append(wav)
  1005→            prefix_texts.append(prefix_text)
  1006→            targets.append(target)
  1007→            sample_debug.append({
  1008→                "sample_id": _as_str(f.get("sample_id", "")),
  1009→                "tar_member_name": _as_str(f.get("tar_member_name", "")),
  1010→                "bucket_id": _as_str(f.get("bucket_id", "")),
  1011→                "duration_s": round(float(duration_s), 3),
  1012→                "transcript_len": len(transcript),
  1013→            })
  1014→
  1015→        if not audios:
  1016→            import warnings
  1017→            warnings.warn("All samples in batch failed decoding/filtering — returning dummy batch")
  1018→            return self._make_zero_loss_batch(
  1019→                batch_debug={"sample_count": 0, "samples": sample_debug},
  1020→                skip_reason="all_samples_filtered",
  1021→            )
  1022→
  1023→        eos = self.processor.tokenizer.eos_token or ""
  1024→        full_texts = [pfx + tgt + eos for pfx, tgt in zip(prefix_texts, targets)]
  1025→
  1026→        full_inputs = self.processor(
  1027→            text=full_texts,
  1028→            audio=audios,
  1029→            return_tensors="pt",
  1030→            padding=True,
  1031→            truncation=False,
  1032→        )
  1033→
  1034→        audio_token_lens = _get_feat_extract_output_lengths(
  1035→            full_inputs["feature_attention_mask"].sum(dim=1)
  1036→        ).tolist()
  1037→        prefix_expanded = self.processor.replace_multimodal_special_tokens(
  1038→            prefix_texts,
  1039→            iter(audio_token_lens),
  1040→        )
  1041→        prefix_tok = self.processor.tokenizer(
  1042→            prefix_expanded,
  1043→            return_tensors="pt",
  1044→            padding=True,
  1045→            truncation=False,
  1046→        )
  1047→        prefix_lens = prefix_tok["attention_mask"].sum(dim=1).tolist()
  1048→
  1049→        labels = full_inputs["input_ids"].clone()
  1050→        for i, pl in enumerate(prefix_lens):
  1051→            valid_token_positions = torch.nonzero(full_inputs["attention_mask"][i], as_tuple=True)[0]
  1052→            labels[i, valid_token_positions[:pl]] = -100
  1053→
  1054→        pad_id = self.processor.tokenizer.pad_token_id
  1055→        if pad_id is not None:
  1056→            labels[labels == pad_id] = -100
  1057→
  1058→        batch_audio_secs = sum(len(a) / self.sampling_rate for a in audios)
  1059→        batch_debug = {
  1060→            "bucket_id": ",".join(sorted({s.get("bucket_id", "") for s in sample_debug if s.get("bucket_id", "")})),


     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→import base64
    17→import io
    18→import urllib.request
    19→from dataclasses import dataclass
    20→from typing import Any, Iterable, List, Optional, Tuple, Union
    21→from urllib.parse import urlparse
    22→
    23→import librosa
    24→import numpy as np
    25→import soundfile as sf
    26→
    27→AudioLike = Union[
    28→    str,                      # wav path / URL / base64
    29→    Tuple[np.ndarray, int],   # (waveform, sr)
    30→]
    31→MaybeList = Union[Any, List[Any]]
    32→
    33→SAMPLE_RATE = 16000
    34→MAX_ASR_INPUT_SECONDS = 1200
    35→MAX_FORCE_ALIGN_INPUT_SECONDS = 180
    36→MIN_ASR_INPUT_SECONDS = 0.5
    37→SUPPORTED_LANGUAGES: List[str] = [
    38→    "Chinese",
    39→    "English",
    40→    "Cantonese",
    41→    "Arabic",
    42→    "German",
    43→    "French",
    44→    "Spanish",
    45→    "Portuguese",
    46→    "Indonesian",
    47→    "Italian",
    48→    "Korean",
    49→    "Russian",
    50→    "Thai",
    51→    "Vietnamese",
    52→    "Japanese",
    53→    "Turkish",
    54→    "Hindi",
    55→    "Malay",
    56→    "Dutch",
    57→    "Swedish",
    58→    "Danish",
    59→    "Finnish",
    60→    "Polish",
    61→    "Czech",
    62→    "Filipino",
    63→    "Persian",
    64→    "Greek",
    65→    "Romanian",
    66→    "Hungarian",
    67→    "Macedonian",
    68→    # Indic languages (added via Phase 2 fine-tuning)
    69→    "Assamese",
    70→    "Bengali",
    71→    "Gujarati",
    72→    "Kannada",
    73→    "Malayalam",
    74→    "Marathi",
    75→    "Odia",
    76→    "Punjabi",
    77→    "Tamil",
    78→    "Telugu",
    79→]
    80→_ASR_TEXT_TAG = "<asr_text>"
    81→_LANG_PREFIX = "language "
    82→
    83→
    84→def normalize_language_name(language: str) -> str:
    85→    """
    86→    Normalize language name to the canonical format used by Qwen3-ASR:
    87→    first letter uppercase, the rest lowercase (e.g., 'cHINese' -> 'Chinese').
    88→
    89→    Args:
    90→        language (str): Input language name.
    91→
    92→    Returns:
    93→        str: Normalized language name.
    94→
    95→    Raises:
    96→        ValueError: If language is empty.
    97→    """
    98→    if language is None:
    99→        raise ValueError("language is None")
   100→    s = str(language).strip()
   101→    if not s:
   102→        raise ValueError("language is empty")
   103→    return s[:1].upper() + s[1:].lower()
   104→
   105→
   106→def validate_language(language: str) -> None:
   107→    """
   108→    Validate the language is supported.
   109→
   110→    Args:
   111→        language (str): Canonical language name.
   112→
   113→    Raises:
   114→        ValueError: If unsupported.
   115→    """
   116→    if language not in SUPPORTED_LANGUAGES:
   117→        raise ValueError(f"Unsupported language: {language}. Supported: {SUPPORTED_LANGUAGES}")
   118→
   119→
   120→def ensure_list(x: MaybeList) -> List[Any]:
   121→        return x if isinstance(x, list) else [x]
   122→
   123→
   124→def is_url(s: str) -> bool:
   125→    try:
   126→        u = urlparse(s)
   127→        return u.scheme in ("http", "https") and bool(u.netloc)
   128→    except Exception:
   129→        return False
   130→
   131→
   132→def is_probably_base64(s: str) -> bool:
   133→    if s.startswith("data:audio"):
   134→        return True
   135→    if ("/" not in s and "\\" not in s) and len(s) > 256:
   136→        return True
   137→    return False
   138→
   139→
   140→def decode_base64_bytes(b64: str) -> bytes:
   141→    if "," in b64 and b64.strip().startswith("data:"):
   142→        b64 = b64.split(",", 1)[1]
   143→    return base64.b64decode(b64)
   144→
   145→
   146→def load_audio_any(x: str) -> Tuple[np.ndarray, int]:
   147→    if is_url(x):
   148→        with urllib.request.urlopen(x) as resp:
   149→            audio_bytes = resp.read()
   150→        with io.BytesIO(audio_bytes) as f:
   151→            audio, sr = sf.read(f, dtype="float32", always_2d=False)
   152→    elif is_probably_base64(x):
   153→        audio_bytes = decode_base64_bytes(x)
   154→        with io.BytesIO(audio_bytes) as f:
   155→            audio, sr = sf.read(f, dtype="float32", always_2d=False)
   156→    else:
   157→        audio, sr = librosa.load(x, sr=None, mono=False)
   158→
   159→    audio = np.asarray(audio, dtype=np.float32)
   160→    sr = int(sr)
   161→    return audio, sr
   162→
   163→
   164→def to_mono(audio: np.ndarray) -> np.ndarray:
   165→    if audio.ndim == 1:
   166→        return audio
   167→    # soundfile can return shape (T, C); some pipelines use (C, T)
   168→    if audio.ndim == 2:
   169→        if audio.shape[0] <= 8 and audio.shape[1] > audio.shape[0]:
   170→            audio = audio.T
   171→        return np.mean(audio, axis=-1).astype(np.float32)
   172→    raise ValueError(f"Unsupported audio ndim={audio.ndim}")
   173→
   174→
   175→def float_range_normalize(audio: np.ndarray) -> np.ndarray:
   176→    audio = audio.astype(np.float32)
   177→    if audio.size == 0:
   178→        return audio
   179→    peak = float(np.max(np.abs(audio)))
   180→    if peak == 0.0:
   181→        return audio
   182→    # If decoded audio is int-like scaled or out-of-range, normalize conservatively.
   183→    if peak > 1.0:
   184→        audio = audio / peak
   185→    audio = np.clip(audio, -1.0, 1.0)
   186→    return audio
   187→
   188→
   189→def normalize_audio_input(a: AudioLike) -> np.ndarray:
   190→    """
   191→    Normalize one audio input to mono 16k float32 waveform in [-1, 1].
   192→
   193→    Supported inputs:
   194→        - str: local file path / https URL / base64 audio string
   195→        - (np.ndarray, sr): waveform and sampling rate
   196→
   197→    Returns:
   198→        np.ndarray:
   199→            Mono 16k float32 waveform in [-1, 1].
   200→    """
   201→    if isinstance(a, str):
   202→        audio, sr = load_audio_any(a)
   203→    elif isinstance(a, tuple) and len(a) == 2 and isinstance(a[0], np.ndarray):
   204→        audio, sr = a[0], int(a[1])
   205→    else:
   206→        raise TypeError(f"Unsupported audio input type: {type(a)}")
   207→
   208→    audio = to_mono(np.asarray(audio))
   209→    if sr != SAMPLE_RATE:
   210→        audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE).astype(np.float32)
   211→    audio = float_range_normalize(audio)
   212→    return audio
   213→
   214→
   215→def normalize_audios(audios: Union[AudioLike, List[AudioLike]]) -> List[np.ndarray]:
   216→    items = ensure_list(audios)
   217→    return [normalize_audio_input(a) for a in items]
   218→
   219→
   220→def chunk_list(xs: List[Any], chunk_size: int) -> Iterable[List[Any]]:
   221→    """
   222→    Yield chunks of a list.
   223→
   224→    Args:
   225→        xs (List[Any]): Input list.
   226→        chunk_size (int): Chunk size.
   227→
   228→    Yields:
   229→        List[Any]: Slices of xs.
   230→    """
   231→    if chunk_size <= 0:
   232→        yield xs
   233→        return
   234→    for i in range(0, len(xs), chunk_size):
   235→        yield xs[i : i + chunk_size]
   236→
   237→
   238→@dataclass(frozen=True)
   239→class AudioChunk:
   240→    """
   241→    One chunk cut from an original audio.
   242→
   243→    Attributes:
   244→        orig_index: Index of the original sample in the input batch.
   245→        chunk_index: Index of this chunk within the original sample.
   246→        wav: Mono float32 waveform.
   247→        sr: Sampling rate.
   248→        offset_sec: Start offset of this chunk in the original audio, in seconds.
   249→    """
   250→    orig_index: int
   251→    chunk_index: int
   252→    wav: np.ndarray
   253→    sr: int
   254→    offset_sec: float
   255→
   256→
   257→def split_audio_into_chunks(
   258→    wav: np.ndarray,
   259→    sr: int,
   260→    max_chunk_sec: float,
   261→    search_expand_sec: float = 5.0,
   262→    min_window_ms: float = 100.0,
   263→) -> List[Tuple[np.ndarray, float]]:
   264→    """
   265→    Split a long audio into chunks close to max_chunk_sec, using a low-energy boundary.
   266→
   267→    This implementation guarantees:
   268→      - Concatenating all returned chunks reproduces the original audio exactly
   269→        (total number of samples is identical, no overlaps, no gaps).
   270→
   271→    Args:
   272→        wav: Mono waveform float32.
   273→        sr: Sampling rate.
   274→        max_chunk_sec: Target max chunk duration in seconds.
   275→        search_expand_sec: Boundary search half-window in seconds.
   276→        min_window_ms: Sliding window in milliseconds for energy estimation.
   277→
   278→    Returns:
   279→        List[Tuple[np.ndarray, float]]: List of (chunk_wav, offset_sec).
   280→    """
   281→    wav = np.asarray(wav, dtype=np.float32)
   282→    if wav.ndim > 1:
   283→        wav = np.mean(wav, axis=-1).astype(np.float32)
   284→
   285→    total_len = int(wav.shape[0])
   286→    total_sec = total_len / float(sr)
   287→    if total_sec <= max_chunk_sec:
   288→        return [(wav, 0.0)]
   289→
   290→    max_len = int(max_chunk_sec * sr)
   291→    expand = int(search_expand_sec * sr)
   292→    win = max(4, int((min_window_ms / 1000.0) * sr))
   293→
   294→    chunks: List[Tuple[np.ndarray, float]] = []
   295→
   296→    start = 0
   297→    offset_sec = 0.0
   298→
   299→    while (total_len - start) > max_len:
   300→        cut = start + max_len
   301→
   302→        left = max(start, cut - expand)
   303→        right = min(total_len, cut + expand)
   304→
   305→        if right - left <= win:
   306→            boundary = cut
   307→        else:
   308→            seg = wav[left:right]
   309→            seg_abs = np.abs(seg)
   310→
   311→            window_sums = np.convolve(seg_abs, np.ones(win, dtype=np.float32), mode="valid")
   312→
   313→            min_pos = int(np.argmin(window_sums))
   314→
   315→            wstart = min_pos
   316→            wend = min_pos + win
   317→            local = seg_abs[wstart:wend]
   318→            inner = int(np.argmin(local))
   319→            boundary = left + wstart + inner
   320→
   321→        boundary = int(max(boundary, start + 1))
   322→        boundary = int(min(boundary, total_len))
   323→
   324→        chunk = wav[start:boundary]
   325→        chunks.append((chunk, offset_sec))
   326→
   327→        offset_sec += (boundary - start) / float(sr)
   328→        start = boundary
   329→
   330→    tail = wav[start:total_len]
   331→    chunks.append((tail, offset_sec))
   332→
   333→    # Pad too-short chunks to at least MIN_ASR_INPUT_SECONDS (zero-padding at tail)
   334→    min_len = int(MIN_ASR_INPUT_SECONDS * sr)
   335→    padded: List[Tuple[np.ndarray, float]] = []
   336→    for c, off in chunks:
   337→        if c.shape[0] < min_len:
   338→            pad = min_len - int(c.shape[0])
   339→            c = np.pad(c, (0, pad), mode="constant", constant_values=0.0).astype(np.float32)
   340→        padded.append((c, off))
   341→    chunks = padded
   342→
   343→    return chunks
   344→
   345→
   346→def detect_and_fix_repetitions(text, threshold=20):
   347→    def fix_char_repeats(s, thresh):
   348→        res = []
   349→        i = 0
   350→        n = len(s)
   351→        while i < n:
   352→            count = 1
   353→            while i + count < n and s[i + count] == s[i]:
   354→                count += 1
   355→
   356→            if count > thresh:
   357→                res.append(s[i])
   358→                i += count
   359→            else:
   360→                res.append(s[i:i+count])
   361→                i += count
   362→        return ''.join(res)
   363→
   364→    def fix_pattern_repeats(s, thresh, max_len=20):
   365→        n = len(s)
   366→        min_repeat_chars = thresh * 2
   367→        if n < min_repeat_chars:
   368→            return s
   369→
   370→        i = 0
   371→        result = []
   372→        while i <= n - min_repeat_chars:
   373→            found = False
   374→            for k in range(1, max_len + 1):
   375→                if i + k * thresh > n:
   376→                    break
   377→
   378→                pattern = s[i:i+k]
   379→                valid = True
   380→                for rep in range(1, thresh):
   381→                    start_idx = i + rep * k
   382→                    if s[start_idx:start_idx+k] != pattern:
   383→                        valid = False
   384→                        break
   385→
   386→                if valid:
   387→                    total_rep = thresh
   388→                    end_index = i + thresh * k
   389→                    while end_index + k <= n and s[end_index:end_index+k] == pattern:
   390→                        total_rep += 1
   391→                        end_index += k
   392→                    result.append(pattern)
   393→                    result.append(fix_pattern_repeats(s[end_index:], thresh, max_len))
   394→                    i = n
   395→                    found = True
   396→                    break
   397→
   398→            if found:
   399→                break
   400→            else:
   401→                result.append(s[i])
   402→                i += 1
   403→
   404→        if not found:
   405→            result.append(s[i:])
   406→        return ''.join(result)
   407→
   408→    text_raw = text
   409→    text = fix_char_repeats(text_raw, threshold)
   410→    text = fix_pattern_repeats(text, threshold)
   411→    return text
   412→
   413→
   414→def parse_asr_output(
   415→    raw: str,
   416→    user_language: Optional[str] = None,
   417→) -> Tuple[str, str]:
   418→    """
   419→    Parse Qwen3-ASR raw output into (language, text).
   420→
   421→    Cases:
   422→      - With tag: "language Chinese<asr_text>...."
   423→      - With newlines: "language Chinese\\n...\\n<asr_text>...."
   424→      - No tag: treat whole string as text.
   425→      - "language None<asr_text>": treat as empty audio -> ("", "")
   426→
   427→    If user_language is provided, language is forced to user_language and raw is treated as text-only
   428→    (the model is expected to output plain transcription without metadata).
   429→
   430→    Args:
   431→        raw: Raw decoded string.
   432→        user_language: Canonical language name if user forced language.
   433→
   434→    Returns:
   435→        Tuple[str, str]: (language, text)
   436→    """
   437→    if raw is None:
   438→        return "", ""
   439→    s = str(raw).strip()
   440→    if not s:
   441→        return "", ""
   442→
   443→    s = detect_and_fix_repetitions(s)
   444→
   445→    if user_language:
   446→        # user explicitly forced language => model output is treated as pure text
   447→        return user_language, s
   448→
   449→    meta_part = s
   450→    text_part = ""
   451→    has_tag = _ASR_TEXT_TAG in s
   452→    if has_tag:
   453→        meta_part, text_part = s.split(_ASR_TEXT_TAG, 1)
   454→    else:
   455→        # no tag => pure text
   456→        return "", s.strip()
   457→
   458→    meta_lower = meta_part.lower()
   459→
   460→    # empty audio heuristic
   461→    if "language none" in meta_lower:
   462→        t = text_part.strip()
   463→        if not t:
   464→            return "", ""
   465→        # if model still returned something, keep it but language unknown
   466→        return "", t
   467→
   468→    # extract "language xxx" from meta
   469→    lang = ""
   470→    for line in meta_part.splitlines():
   471→        line = line.strip()
   472→        if not line:
   473→            continue
   474→        low = line.lower()
   475→        if low.startswith(_LANG_PREFIX):
   476→            val = line[len(_LANG_PREFIX):].strip()
   477→            if val:
   478→                lang = normalize_language_name(val)
   479→            break
   480→
   481→    return lang, text_part.strip()
   482→
   483→
   484→def merge_languages(langs: List[str]) -> str:
   485→    """
   486→    Merge per-chunk languages into a compact comma-separated string,
   487→    keeping order and removing consecutive duplicates and empty entries.
   488→
   489→    Example:
   490→      ["Chinese", "English", "English"] -> "Chinese,English"
   491→
   492→    Args:
   493→        langs: List of canonical language names.
   494→
   495→    Returns:
   496→        str: Merged language string.
   497→    """
   498→    out: List[str] = []
   499→    prev = None
   500→    for x in langs:
   501→        x = (x or "").strip()
   502→        if not x:
   503→            continue
   504→        if x == prev:
   505→            continue
   506→        out.append(x)
   507→        prev = x
   508→    return ",".join(out)
   509→


   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→


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→"""
    17→A gradio demo for Qwen3 ASR models.
    18→"""
    19→
    20→import argparse
    21→import base64
    22→import io
    23→import json
    24→import os
    25→from typing import Any, Dict, List, Optional, Tuple, Union
    26→
    27→import gradio as gr
    28→import numpy as np
    29→import torch
    30→from qwen_asr import Qwen3ASRModel
    31→from scipy.io.wavfile import write as wav_write
    32→
    33→
    34→def _title_case_display(s: str) -> str:
    35→    s = (s or "").strip()
    36→    s = s.replace("_", " ")
    37→    return " ".join([w[:1].upper() + w[1:] if w else "" for w in s.split()])
    38→
    39→
    40→def _build_choices_and_map(items: Optional[List[str]]) -> Tuple[List[str], Dict[str, str]]:
    41→    if not items:
    42→        return [], {}
    43→    display = [_title_case_display(x) for x in items]
    44→    mapping = {d: r for d, r in zip(display, items)}
    45→    return display, mapping
    46→
    47→
    48→def _dtype_from_str(s: str) -> torch.dtype:
    49→    s = (s or "").strip().lower()
    50→    if s in ("bf16", "bfloat16"):
    51→        return torch.bfloat16
    52→    if s in ("fp16", "float16", "half"):
    53→        return torch.float16
    54→    if s in ("fp32", "float32"):
    55→        return torch.float32
    56→    raise ValueError(f"Unsupported torch dtype: {s}. Use bfloat16/float16/float32.")
    57→
    58→
    59→def _normalize_audio(wav, eps=1e-12, clip=True):
    60→    x = np.asarray(wav)
    61→
    62→    if np.issubdtype(x.dtype, np.integer):
    63→        info = np.iinfo(x.dtype)
    64→        if info.min < 0:
    65→            y = x.astype(np.float32) / max(abs(info.min), info.max)
    66→        else:
    67→            mid = (info.max + 1) / 2.0
    68→            y = (x.astype(np.float32) - mid) / mid
    69→    elif np.issubdtype(x.dtype, np.floating):
    70→        y = x.astype(np.float32)
    71→        m = np.max(np.abs(y)) if y.size else 0.0
    72→        if m > 1.0 + 1e-6:
    73→            y = y / (m + eps)
    74→    else:
    75→        raise TypeError(f"Unsupported dtype: {x.dtype}")
    76→
    77→    if clip:
    78→        y = np.clip(y, -1.0, 1.0)
    79→
    80→    if y.ndim > 1:
    81→        y = np.mean(y, axis=-1).astype(np.float32)
    82→
    83→    return y
    84→
    85→
    86→def _audio_to_tuple(audio: Any) -> Optional[Tuple[np.ndarray, int]]:
    87→    """
    88→    Accept gradio audio:
    89→      - {"sampling_rate": int, "data": np.ndarray}
    90→      - (sr, np.ndarray)  [some gradio versions]
    91→    Return: (wav_float32_mono, sr)
    92→    """
    93→    if audio is None:
    94→        return None
    95→
    96→    if isinstance(audio, dict) and "sampling_rate" in audio and "data" in audio:
    97→        sr = int(audio["sampling_rate"])
    98→        wav = _normalize_audio(audio["data"])
    99→        return wav, sr
   100→
   101→    if isinstance(audio, tuple) and len(audio) == 2:
   102→        a0, a1 = audio
   103→        if isinstance(a0, int):
   104→            sr = int(a0)
   105→            wav = _normalize_audio(a1)
   106→            return wav, sr
   107→        if isinstance(a1, int):
   108→            wav = _normalize_audio(a0)
   109→            sr = int(a1)
   110→            return wav, sr
   111→
   112→    return None
   113→
   114→
   115→def _parse_audio_any(audio: Any) -> Union[str, Tuple[np.ndarray, int]]:
   116→    if audio is None:
   117→        raise ValueError("Audio is required.")
   118→    at = _audio_to_tuple(audio)
   119→    if at is not None:
   120→        return at
   121→    raise ValueError("Unsupported audio input format.")
   122→
   123→
   124→def build_parser() -> argparse.ArgumentParser:
   125→    parser = argparse.ArgumentParser(
   126→        prog="qwen-asr-demo",
   127→        description=(
   128→            "Launch a Gradio demo for Qwen3 ASR models (Transformers / vLLM).\n\n"
   129→            "Examples:\n"
   130→            "  qwen-asr-demo --asr-checkpoint Qwen/Qwen3-ASR-1.7B\n"
   131→            "  qwen-asr-demo --asr-checkpoint Qwen/Qwen3-ASR-1.7B --aligner-checkpoint Qwen/Qwen3-ForcedAligner-0.6B\n"
   132→            "  qwen-asr-demo --backend vllm --cuda-visible-devices 0\n"
   133→            "  qwen-asr-demo --backend transformers --backend-kwargs '{\"device_map\":\"cuda:0\",\"dtype\":\"bfloat16\",\"attn_implementation\":\"flash_attention_2\"}'\n"
   134→            "  qwen-asr-demo --backend vllm --backend-kwargs '{\"gpu_memory_utilization\":0.85}'\n"
   135→        ),
   136→        formatter_class=argparse.RawTextHelpFormatter,
   137→        add_help=True,
   138→    )
   139→
   140→    parser.add_argument("--asr-checkpoint", required=True, help="Qwen3-ASR model checkpoint path or HF repo id.")
   141→    parser.add_argument(
   142→        "--aligner-checkpoint",
   143→        default=None,
   144→        help="Qwen3-ForcedAligner checkpoint path or HF repo id (optional; enables timestamps when provided).",
   145→    )
   146→
   147→    parser.add_argument(
   148→        "--backend",
   149→        default="transformers",
   150→        choices=["transformers", "vllm"],
   151→        help="Backend for ASR model loading (default: transformers).",
   152→    )
   153→
   154→    parser.add_argument(
   155→        "--cuda-visible-devices",
   156→        default="0",
   157→        help=(
   158→            "Set CUDA_VISIBLE_DEVICES for the demo process (default: 0). "
   159→            "Use e.g. '0' or '1'"
   160→        ),
   161→    )
   162→
   163→    parser.add_argument(
   164→        "--backend-kwargs",
   165→        default=None,
   166→        help=(
   167→            "JSON dict for backend-specific kwargs excluding checkpoints.\n"
   168→            "Examples:\n"
   169→            "  transformers: '{\"device_map\":\"cuda:0\",\"dtype\":\"bfloat16\",\"attn_implementation\":\"flash_attention_2\",\"max_inference_batch_size\":32}'\n"
   170→            "  vllm        : '{\"gpu_memory_utilization\":0.8,\"max_inference_batch_size\":32}'\n"
   171→        ),
   172→    )
   173→    parser.add_argument(
   174→        "--aligner-kwargs",
   175→        default=None,
   176→        help=(
   177→            "JSON dict for forced aligner kwargs (only used when --aligner-checkpoint is set).\n"
   178→            "Example: '{\"dtype\":\"bfloat16\",\"device_map\":\"cuda:0\"}'\n"
   179→        ),
   180→    )
   181→
   182→    # Gradio server args
   183→    parser.add_argument("--ip", default="0.0.0.0", help="Server bind IP for Gradio (default: 0.0.0.0).")
   184→    parser.add_argument("--port", type=int, default=8000, help="Server port for Gradio (default: 8000).")
   185→    parser.add_argument(
   186→        "--share/--no-share",
   187→        dest="share",
   188→        default=False,
   189→        action=argparse.BooleanOptionalAction,
   190→        help="Whether to create a public Gradio link (default: disabled).",
   191→    )
   192→    parser.add_argument("--concurrency", type=int, default=16, help="Gradio queue concurrency (default: 16).")
   193→
   194→    # HTTPS args
   195→    parser.add_argument("--ssl-certfile", default=None, help="Path to SSL certificate file for HTTPS (optional).")
   196→    parser.add_argument("--ssl-keyfile", default=None, help="Path to SSL key file for HTTPS (optional).")
   197→    parser.add_argument(
   198→        "--ssl-verify/--no-ssl-verify",
   199→        dest="ssl_verify",
   200→        default=True,
   201→        action=argparse.BooleanOptionalAction,
   202→        help="Whether to verify SSL certificate (default: enabled).",
   203→    )
   204→
   205→    return parser
   206→
   207→
   208→def _parse_json_dict(s: Optional[str], *, name: str) -> Dict[str, Any]:
   209→    if s is None or not str(s).strip():
   210→        return {}
   211→    try:
   212→        obj = json.loads(s)
   213→    except Exception as e:
   214→        raise ValueError(f"Invalid JSON for {name}: {e}")
   215→    if not isinstance(obj, dict):
   216→        raise ValueError(f"{name} must be a JSON object (dict).")
   217→    return obj
   218→
   219→
   220→def _apply_cuda_visible_devices(cuda_visible_devices: str) -> None:
   221→    v = (cuda_visible_devices or "").strip()
   222→    if not v:
   223→        return
   224→    os.environ["CUDA_VISIBLE_DEVICES"] = v
   225→
   226→
   227→def _default_backend_kwargs(backend: str) -> Dict[str, Any]:
   228→    if backend == "transformers":
   229→        return dict(
   230→            dtype=torch.bfloat16,
   231→            device_map="cuda:0",
   232→            max_inference_batch_size=4,
   233→            max_new_tokens=512,
   234→        )
   235→    else:
   236→        return dict(
   237→            gpu_memory_utilization=0.8,
   238→            max_inference_batch_size=4,
   239→            max_new_tokens=4096,
   240→        )
   241→
   242→
   243→def _default_aligner_kwargs() -> Dict[str, Any]:
   244→    return dict(
   245→        dtype=torch.bfloat16,
   246→        device_map="cuda:0",
   247→    )
   248→
   249→
   250→def _merge_dicts(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
   251→    out = dict(base)
   252→    out.update(override)
   253→    return out
   254→
   255→
   256→def _coerce_special_types(d: Dict[str, Any]) -> Dict[str, Any]:
   257→    out: Dict[str, Any] = {}
   258→    for k, v in d.items():
   259→        if k == "dtype" and isinstance(v, str):
   260→            out[k] = _dtype_from_str(v)
   261→        else:
   262→            out[k] = v
   263→    return out
   264→
   265→
   266→def _make_timestamp_html(audio_upload: Any, timestamps: Any) -> str:
   267→    """
   268→    Build HTML with per-token audio slices, using base64 data URLs.
   269→    Expect timestamps as list[dict] with keys: text, start_time, end_time (ms).
   270→    """
   271→    at = _audio_to_tuple(audio_upload)
   272→    if at is None:
   273→        raise ValueError("Audio input is required for visualization.")
   274→    audio, sr = at
   275→
   276→    if not timestamps:
   277→        return "
No timestamps to visualize.
"
   278→    if not isinstance(timestamps, list):
   279→        raise ValueError("Timestamps must be a list (JSON array).")
   280→
   281→    html_content = """
   282→
   295→    """
   296→
   297→    html_content += """
   298→
   299→        Timestamps Visualization (时间戳可视化结果）
   300→

   301→    """
   302→
   303→    for item in timestamps:
   304→        if not isinstance(item, dict):
   305→            continue
   306→        word = str(item.get("text", "") or "")
   307→        start = item.get("start_time", None)
   308→        end = item.get("end_time", None)
   309→        if start is None or end is None:
   310→            continue
   311→
   312→        start = float(start)
   313→        end = float(end)
   314→        if end <= start:
   315→            continue
   316→
   317→        start_sample = max(0, int(start * sr))
   318→        end_sample = min(len(audio), int(end * sr))
   319→        if end_sample <= start_sample:
   320→            continue
   321→
   322→        seg = audio[start_sample:end_sample]
   323→        seg_i16 = (np.clip(seg, -1.0, 1.0) * 32767.0).astype(np.int16)
   324→
   325→        mem = io.BytesIO()
   326→        wav_write(mem, sr, seg_i16)
   327→        mem.seek(0)
   328→        b64 = base64.b64encode(mem.read()).decode("utf-8")
   329→        audio_src = f"data:audio/wav;base64,{b64}"
   330→
   331→        html_content += f"""
   332→

   333→
{word}

   334→
{start} - {end} s

   335→

   336→                <audio controls preload="none" src="{audio_src}"></audio>
   337→

   338→

   339→        """
   340→
   341→    html_content += "
"
   342→    return html_content
   343→
   344→
   345→def build_demo(
   346→    asr: Qwen3ASRModel,
   347→    asr_ckpt: str,
   348→    backend: str,
   349→    aligner_ckpt: Optional[str] = None,
   350→) -> gr.Blocks:
   351→    supported_langs_raw = asr.get_supported_languages()
   352→    lang_choices_disp, lang_map = _build_choices_and_map([x for x in supported_langs_raw])
   353→    lang_choices = ["Auto"] + lang_choices_disp
   354→
   355→    has_aligner = bool(aligner_ckpt)
   356→
   357→    theme = gr.themes.Soft(
   358→        font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"],
   359→    )
   360→    css = ".gradio-container {max-width: none !important;}"
   361→
   362→    with gr.Blocks(theme=theme, css=css) as demo:
   363→        gr.Markdown(
   364→            f"""
   365→# Qwen3 ASR Demo
   366→**Backend:** `{backend}`
   367→**ASR Checkpoint:** `{asr_ckpt}`
   368→**Forced Aligner:** `{aligner_ckpt if aligner_ckpt else "(none)"}`
   369→"""
   370→        )
   371→
   372→        with gr.Row():
   373→            with gr.Column(scale=2):
   374→                audio_in = gr.Audio(label="Audio Input (上传音频)", type="numpy")
   375→                lang_in = gr.Dropdown(
   376→                    label="Language (语种)",
   377→                    choices=lang_choices,
   378→                    value="Auto",
   379→                    interactive=True,
   380→                )
   381→                if has_aligner:
   382→                    ts_in = gr.Checkbox(
   383→                        label="Return Timestamps (是否返回时间戳)",
   384→                        value=True,
   385→                    )
   386→                else:
   387→                    ts_in = gr.State(False)
   388→
   389→                btn = gr.Button("Transcribe (识别)", variant="primary")
   390→
   391→            with gr.Column(scale=2):
   392→                out_lang = gr.Textbox(label="Detected Language", lines=1)
   393→                out_text = gr.Textbox(label="Result Text", lines=12)
   394→
   395→            if has_aligner:
   396→                with gr.Column(scale=3):
   397→                    out_ts = gr.JSON(label="Timestamps（时间戳结果）")
   398→                    viz_btn = gr.Button("Visualize Timestamps (可视化时间戳)", variant="secondary")
   399→            else:
   400→                with gr.Column(scale=3):
   401→                    out_ts = gr.State(None)
   402→                    viz_btn = gr.State(None)
   403→
   404→        # Put the visualization panel below the three columns
   405→        if has_aligner:
   406→            with gr.Row():
   407→                out_ts_html = gr.HTML(label="Timestamps Visualization (时间戳可视化结果)")
   408→        else:
   409→            out_ts_html = gr.State("")
   410→
   411→        def run(audio_upload: Any, lang_disp: str, return_ts: bool):
   412→            audio_obj = _parse_audio_any(audio_upload)
   413→
   414→            language = None
   415→            if lang_disp and lang_disp != "Auto":
   416→                language = lang_map.get(lang_disp, lang_disp)
   417→
   418→            return_ts = bool(return_ts) and has_aligner
   419→
   420→            results = asr.transcribe(
   421→                audio=audio_obj,
   422→                language=language,
   423→                return_time_stamps=return_ts,
   424→            )
   425→            if not isinstance(results, list) or len(results) != 1:
   426→                raise RuntimeError(
   427→                    f"Unexpected result size: {type(results)} "
   428→                    f"len={len(results) if isinstance(results, list) else 'N/A'}"
   429→                )
   430→
   431→            r = results[0]
   432→
   433→            if has_aligner:
   434→                ts_payload = None
   435→                if return_ts:
   436→                    ts_payload = [
   437→                        dict(
   438→                            text=getattr(t, "text", None),
   439→                            start_time=getattr(t, "start_time", None),
   440→                            end_time=getattr(t, "end_time", None),
   441→                        )
   442→                        for t in (getattr(r, "time_stamps", None) or [])
   443→                    ]
   444→                return (
   445→                    getattr(r, "language", "") or "",
   446→                    getattr(r, "text", "") or "",
   447→                    gr.update(value=ts_payload) if return_ts else gr.update(value=None),
   448→                    gr.update(value=""),  # clear html on each transcribe
   449→                )
   450→            else:
   451→                return (
   452→                    getattr(r, "language", "") or "",
   453→                    getattr(r, "text", "") or "",
   454→                )
   455→
   456→        def visualize(audio_upload: Any, timestamps_json: Any):
   457→            return _make_timestamp_html(audio_upload, timestamps_json)
   458→
   459→        if has_aligner:
   460→            btn.click(
   461→                run,
   462→                inputs=[audio_in, lang_in, ts_in],
   463→                outputs=[out_lang, out_text, out_ts, out_ts_html],
   464→            )
   465→            viz_btn.click(
   466→                visualize,
   467→                inputs=[audio_in, out_ts],
   468→                outputs=[out_ts_html],
   469→            )
   470→        else:
   471→            btn.click(
   472→                run,
   473→                inputs=[audio_in, lang_in, ts_in],
   474→                outputs=[out_lang, out_text],
   475→            )
   476→
   477→    return demo
   478→
   479→
   480→def main(argv=None) -> int:
   481→    parser = build_parser()
   482→    args = parser.parse_args(argv)
   483→
   484→    _apply_cuda_visible_devices(args.cuda_visible_devices)
   485→
   486→    backend = args.backend
   487→    asr_ckpt = args.asr_checkpoint
   488→    aligner_ckpt = args.aligner_checkpoint
   489→
   490→    user_backend_kwargs = _parse_json_dict(args.backend_kwargs, name="--backend-kwargs")
   491→    user_aligner_kwargs = _parse_json_dict(args.aligner_kwargs, name="--aligner-kwargs")
   492→
   493→    backend_kwargs = _merge_dicts(_default_backend_kwargs(backend), user_backend_kwargs)
   494→    backend_kwargs = _coerce_special_types(backend_kwargs)
   495→
   496→    forced_aligner = None
   497→    forced_aligner_kwargs = None
   498→    if aligner_ckpt:
   499→        forced_aligner = aligner_ckpt
   500→        aligner_kwargs = _merge_dicts(_default_aligner_kwargs(), user_aligner_kwargs)
   501→        forced_aligner_kwargs = _coerce_special_types(aligner_kwargs)
   502→
   503→    if backend == "transformers":
   504→        asr = Qwen3ASRModel.from_pretrained(
   505→            asr_ckpt,
   506→            forced_aligner=forced_aligner,
   507→            forced_aligner_kwargs=forced_aligner_kwargs,
   508→            **backend_kwargs,
   509→        )
   510→    else:
   511→        asr = Qwen3ASRModel.LLM(
   512→            model=asr_ckpt,
   513→            forced_aligner=forced_aligner,
   514→            forced_aligner_kwargs=forced_aligner_kwargs,
   515→            **backend_kwargs,
   516→        )
   517→
   518→    demo = build_demo(asr, asr_ckpt, backend, aligner_ckpt=aligner_ckpt)
   519→
   520→    launch_kwargs: Dict[str, Any] = dict(
   521→        server_name=args.ip,
   522→        server_port=args.port,
   523→        share=args.share,
   524→        ssl_verify=True if args.ssl_verify else False,
   525→    )
   526→    if args.ssl_certfile is not None:
   527→        launch_kwargs["ssl_certfile"] = args.ssl_certfile
   528→    if args.ssl_keyfile is not None:
   529→        launch_kwargs["ssl_keyfile"] = args.ssl_keyfile
   530→
   531→    demo.queue(default_concurrency_limit=int(args.concurrency)).launch(**launch_kwargs)
   532→    return 0
   533→
   534→
   535→if __name__ == "__main__":
   536→    raise SystemExit(main())
   537→


grep -n "class.*Audio\|def forward" /home/ubuntu/training/qwen3-asr-1.7b-phase2-sft/qwen_asr/core/transformers_backend/modeling_qwen3_asr.py | head -30


60:    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
176:    def forward(
231:    def forward(self, x):
248:    def forward(
451:class Qwen3ASRAudioAttention(nn.Module):
477:    def forward(
523:class Qwen3ASRAudioEncoderLayer(GradientCheckpointingLayer):
536:    def forward(
593:    def forward(self, seqlen: int):
603:class Qwen3ASRAudioEncoder(Qwen3ASRPreTrainedModel):
669:    def forward(
821:    def forward(self, x, position_ids):
851:    def forward(self, x):
866:    def forward(self, hidden_states):
911:    def forward(
988:    def forward(
1159:    def forward(


   603→class Qwen3ASRAudioEncoder(Qwen3ASRPreTrainedModel):
   604→    config: Qwen3ASRAudioEncoderConfig
   605→    main_input_name = "input_features"
   606→    _no_split_modules = ["Qwen3ASRAudioEncoderLayer"]
   607→    _supports_sdpa = True
   608→
   609→    def __init__(self, config: Qwen3ASRAudioEncoderConfig):
   610→        super().__init__(config)
   611→        self.dropout = config.dropout
   612→
   613→        embed_dim = config.d_model
   614→        self.num_mel_bins = config.num_mel_bins
   615→        self.max_source_positions = config.max_source_positions
   616→        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
   617→        self.n_window = config.n_window
   618→        self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim)
   619→        self.layers = nn.ModuleList([Qwen3ASRAudioEncoderLayer(config) for _ in range(config.encoder_layers)])
   620→        self.ln_post = nn.LayerNorm(config.d_model)
   621→        self.gradient_checkpointing = False
   622→        self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
   623→        self.conv2d2 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
   624→        self.conv2d3 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
   625→        self.conv_out = nn.Linear(
   626→            config.downsample_hidden_size * ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2),
   627→            config.d_model,
   628→            bias=False,
   629→        )
   630→        self.proj1 = nn.Linear(config.d_model, config.d_model)
   631→        self.act = ACT2FN[config.activation_function]
   632→        self.proj2 = nn.Linear(config.d_model, config.output_dim)
   633→        self.n_window_infer = self.config.n_window_infer
   634→        self.conv_chunksize = self.config.conv_chunksize
   635→        # Initialize weights and apply final processing
   636→        self.post_init()
   637→
   638→    def _freeze_parameters(self):
   639→        for param in self.parameters():
   640→            param.requires_grad = False
   641→        self._requires_grad = False
   642→
   643→    def get_input_embeddings(self) -> nn.Module:
   644→        return self.conv1
   645→
   646→    def set_input_embeddings(self, value: nn.Module):
   647→        self.conv1 = value
   648→
   649→    def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
   650→        # Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
   651→        # NOTE: the created attention masl only approximates the ragged FA2 attention by
   652→        # allowing bidirectional attention within `cu_seqlens` blocks, and not attending between
   653→        # blocks. Though it will not be a 100% match for FA2's `varlen` path
   654→        if self.config._attn_implementation == "flash_attention_2":
   655→            return None
   656→
   657→        seq_length = inputs_tensor.shape[0]
   658→        attention_mask = torch.full(
   659→            [1, 1, seq_length, seq_length],
   660→            torch.finfo(inputs_tensor.dtype).min,
   661→            device=inputs_tensor.device,
   662→            dtype=inputs_tensor.dtype,
   663→        )
   664→        for i in range(1, len(cu_seqlens)):
   665→            attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
   666→        return attention_mask
   667→
   668→    @auto_docstring
   669→    def forward(
   670→        self,
   671→        input_features,
   672→        feature_lens=None,
   673→        aftercnn_lens=None,
   674→    ):
   675→        r"""
   676→        feature_lens (`torch.LongTensor` of shape `(batch_size,)`):
   677→            mel length
   678→        aftercnn_lens (`torch.LongTensor` of shape `(batch_size,)`):
   679→            mel length after cnn
   680→        """
   681→        aftercnn_lens = _get_feat_extract_output_lengths(feature_lens)
   682→        chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
   683→
   684→        chunk_lengths = torch.tensor(
   685→            [self.n_window * 2] * chunk_num.sum(),
   686→            dtype=torch.long,
   687→            device=feature_lens.device,
   688→        )
   689→        tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
   690→        chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
   691→        chunk_lengths[chunk_lengths == 0] = self.n_window * 2
   692→
   693→        chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
   694→        padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
   695→        feature_lens_after_cnn = _get_feat_extract_output_lengths(chunk_lengths)
   696→        padded_mask_after_cnn = nn.utils.rnn.pad_sequence(
   697→            [torch.ones(length, dtype=torch.bool, device=padded_feature.device) for length in feature_lens_after_cnn],
   698→            batch_first=True,
   699→        )
   700→        padded_feature = padded_feature.unsqueeze(1)
   701→        # Split to chunk to avoid OOM during convolution
   702→        padded_embeds = []
   703→        for chunk in padded_feature.split(self.conv_chunksize, dim=0):
   704→            padded_embed = F.gelu(self.conv2d1(chunk))
   705→            padded_embed = F.gelu(self.conv2d2(padded_embed))
   706→            padded_embed = F.gelu(self.conv2d3(padded_embed))
   707→            padded_embeds.append(padded_embed)
   708→        padded_embed = torch.cat(padded_embeds, dim=0)
   709→        b, c, f, t = padded_embed.size()
   710→        padded_embed = self.conv_out(padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f))
   711→
   712→        positional_embedding = (
   713→            self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
   714→            .unsqueeze(0)
   715→            .to(padded_embed.dtype)
   716→        )
   717→        padded_embed = padded_embed + positional_embedding
   718→        hidden_states = padded_embed[padded_mask_after_cnn]
   719→        cu_chunk_lens = [0]
   720→        window_aftercnn = padded_mask_after_cnn.shape[-1] * (self.n_window_infer // (self.n_window * 2))
   721→        for cnn_len in aftercnn_lens:
   722→            cu_chunk_lens += [window_aftercnn] * (cnn_len // window_aftercnn)
   723→            remainder = cnn_len % window_aftercnn
   724→            if remainder != 0:
   725→                cu_chunk_lens += [remainder]
   726→        cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(-1, dtype=torch.int32)
   727→
   728→        for encoder_layer in self.layers:
   729→            layer_outputs = encoder_layer(
   730→                hidden_states,
   731→                cu_seqlens,
   732→            )
   733→
   734→            hidden_states = layer_outputs[0]
   735→
   736→        hidden_states = self.ln_post(hidden_states)
   737→        hidden_states = self.proj1(hidden_states)
   738→        hidden_states = self.act(hidden_states)
   739→        hidden_states = self.proj2(hidden_states)
   740→        return BaseModelOutput(last_hidden_state=hidden_states)
   741→
   742→    def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"):
   743→        """
   744→        Pads a sequence of tensors to their maximum length on indicated `padding_side`.
   745→        Then prepares a mask so that pad tokens are not attended to.
   746→        """
   747→        max_len = tensor_len.max()
   748→        dim = tensor_list[0].shape[0]
   749→        padded_tensor = torch.full(
   750→            size=(len(tensor_list), dim, max_len),
   751→            fill_value=padding_value,
   752→            dtype=self.dtype,
   753→            device=tensor_list[0].device,
   754→        )
   755→
   756→        batch_mask = torch.zeros(
   757→            (len(tensor_len), max_len),
   758→            dtype=torch.long,
   759→            device=padded_tensor.device,
   760→        )
   761→        for i, length in enumerate(tensor_len):
   762→            batch_mask[i, :length] = 1
   763→            padded_tensor[i, :, :length] = tensor_list[i]
   764→
   765→        feature_lens_after_cnn = (tensor_len - 1) // 2 + 1
   766→        max_len_after_cnn = feature_lens_after_cnn.max()
   767→        batch_mask_after_cnn = torch.zeros(
   768→            (len(tensor_len), max_len_after_cnn),
   769→            dtype=torch.long,
   770→            device=padded_tensor.device,
   771→        )
   772→        for i, length in enumerate(feature_lens_after_cnn):
   773→            batch_mask_after_cnn[i, :length] = 1
   774→        return (
   775→            padded_tensor,
   776→            batch_mask.unsqueeze(1),
   777→            batch_mask_after_cnn.bool(),
   778→        )
   779→
   780→
   781→class Qwen3ASRThinkerTextRotaryEmbedding(nn.Module):
   782→    inv_freq: torch.Tensor  # fix linting for `register_buffer`
   783→
   784→    def __init__(self, config: Qwen3ASRConfig, device=None):
   785→        super().__init__()
   786→        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
   787→            self.rope_type = config.rope_scaling.get("rope_type", "default")
   788→        else:
   789→            self.rope_type = "default"
   790→        self.max_seq_len_cached = config.max_position_embeddings
   791→        self.original_max_seq_len = config.max_position_embeddings
   792→
   793→        self.config = config
   794→        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
   795→
   796→        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
   797→        self.register_buffer("inv_freq", inv_freq, persistent=False)
   798→        self.original_inv_freq = self.inv_freq
   799→
   800→        self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
   801→
   802→    def apply_interleaved_mrope(self, freqs, mrope_section):