dnnai/asr_agent/whisper_asr/qwen_engine.py

import numpy as np
from typing import Optional


class QwenASREngine:
    def __init__(
        self,
        model_id: str = "Qwen/Qwen3-ASR-0.6B",
        device: str = "auto",
        language: Optional[str] = None,
    ):
        import torch
        from qwen_asr import Qwen3ASRModel

        self.model_id = model_id
        self.language = language
        self.sample_rate = 16000

        if device == "auto":
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        dtype = torch.float16 if device == "cuda" else torch.float32

        print(f"Loading Qwen3-ASR: {model_id} on {device}")
        self.model = Qwen3ASRModel.from_pretrained(
            model_id, dtype=dtype, device_map=device,
        )
        print("ASR model ready")

    def transcribe_array(self, audio_array: np.ndarray, sample_rate: int = 16000) -> dict:
        if sample_rate != self.sample_rate:
            audio_array = self._resample(audio_array, sample_rate, self.sample_rate)
        result = self.model.transcribe(
            audio=(audio_array, self.sample_rate),
            language=self.language,
        )
        return {"text": result[0].text, "language": result[0].language}

    def _resample(self, audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
        if orig_sr == target_sr:
            return audio
        duration = len(audio) / orig_sr
        new_length = int(duration * target_sr)
        indices = np.linspace(0, len(audio) - 1, new_length)
        return np.interp(indices, np.arange(len(audio)), audio)

    def get_model_info(self) -> dict:
        return {"model_id": self.model_id, "device": self.device}