Spaces:
Sleeping
Sleeping
| import io | |
| import numpy as np | |
| import torch | |
| from pydub import AudioSegment | |
| from transformers import ( | |
| pipeline, | |
| WhisperProcessor, | |
| WhisperForConditionalGeneration, | |
| ) | |
| from typing import Tuple | |
| from app.config import config | |
| class WhisperASR: | |
| """Explicit Whisper loader — avoids pipeline preprocessor num_frames bug.""" | |
| def __init__(self, model_id: str, token: str, device: torch.device): | |
| self.processor = WhisperProcessor.from_pretrained( | |
| model_id, | |
| token=token, | |
| ) | |
| self.model = WhisperForConditionalGeneration.from_pretrained( | |
| model_id, | |
| token=token, | |
| use_safetensors=False, | |
| ).to(device) | |
| self.model.eval() | |
| self.device = device | |
| def transcribe(self, samples: np.ndarray, sampling_rate: int) -> str: | |
| inputs = self.processor( | |
| samples, | |
| sampling_rate=sampling_rate, | |
| return_tensors="pt", | |
| ).to(self.device) | |
| with torch.no_grad(): | |
| predicted_ids = self.model.generate(**inputs) | |
| transcription = self.processor.batch_decode( | |
| predicted_ids, | |
| skip_special_tokens=True, | |
| ) | |
| return transcription[0].strip() | |
| class STTPipeline: | |
| def __init__(self): | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| token = config.HF_TOKEN or None | |
| # Language identifier — wav2vec2, pipeline is fine here | |
| self.lid = pipeline( | |
| "audio-classification", | |
| model=config.LID_MODEL, | |
| device=0 if self.device.type == "cuda" else -1, | |
| token=token, | |
| ) | |
| # ASR models — explicit Whisper loader per language | |
| self.asr: dict = { | |
| lang: WhisperASR(model_id, token, self.device) | |
| for lang, model_id in config.ASR_MODELS.items() | |
| } | |
| def decode_audio(self, audio_bytes: bytes) -> Tuple[np.ndarray, float]: | |
| seg = AudioSegment.from_file(io.BytesIO(audio_bytes)) | |
| seg = seg.set_channels(1).set_frame_rate(config.SAMPLING_RATE) | |
| samples = np.array(seg.get_array_of_samples()).astype(np.float32) | |
| samples /= np.iinfo(seg.array_type).max | |
| duration = len(seg) / 1000.0 | |
| return samples, duration | |
| def classify_language(self, samples: np.ndarray) -> Tuple[str, float]: | |
| result = self.lid( | |
| {"array": samples, "sampling_rate": config.SAMPLING_RATE}, | |
| top_k=1, | |
| ) | |
| label = result[0]["label"].lower() | |
| confidence = result[0]["score"] | |
| return label, confidence | |
| def transcribe(self, audio_bytes: bytes) -> dict: | |
| samples, duration = self.decode_audio(audio_bytes) | |
| language, confidence = self.classify_language(samples) | |
| asr_model = self.asr.get(language) | |
| if asr_model is None: | |
| raise ValueError( | |
| f"No ASR model available for detected language: '{language}'. " | |
| f"Supported languages: {list(self.asr.keys())}" | |
| ) | |
| transcription = asr_model.transcribe(samples, config.SAMPLING_RATE) | |
| return { | |
| "transcription": transcription, | |
| "language": language, | |
| "confidence": round(confidence, 4), | |
| "duration_sec": round(duration, 2), | |
| } | |
| stt_pipeline = STTPipeline() |