Spaces:
Sleeping
Sleeping
File size: 5,551 Bytes
0db822c 9308938 0db822c 54c1995 9308938 54c1995 0db822c 9308938 0db822c 9308938 0db822c 9308938 0db822c 9308938 0db822c 54c1995 0db822c 54c1995 0db822c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import logging
import torch
import torchaudio
import numpy as np
from pathlib import Path
from typing import List, Dict, Any, Optional, Union
# --- Constants ---
TARGET_SR = 16000
WINDOW_SIZE_SEC = 30 # Whisper's native window
VAD_THRESHOLD = 0.50
logger = logging.getLogger(__name__)
class WhisperTranscriber:
def __init__(self, model_path: str = "openai/whisper-large-v3", device: Optional[str] = None):
"""
Initialize the high-performance Whisper pipeline.
Optimized for Egyptian Arabic real estate calls.
"""
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
logger.info("Loading model from %s on %s", model_path, self.device)
try:
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
self.processor = AutoProcessor.from_pretrained(model_path)
dtype = torch.float16 if "cuda" in self.device else torch.float32
self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_path,
dtype=dtype,
low_cpu_mem_usage=True,
).to(self.device)
# batch_size=16 is only useful on GPU; CPU benefits from 1-2 chunks at a time
batch_size = 8 if "cuda" in self.device else 2
self.pipe = pipeline(
"automatic-speech-recognition",
model=self.model,
tokenizer=self.processor.tokenizer,
feature_extractor=self.processor.feature_extractor,
chunk_length_s=30,
batch_size=batch_size,
return_timestamps=True,
dtype=dtype,
device=self.device,
generate_kwargs={"max_new_tokens": 128},
)
except Exception as e:
logger.error("Failed to load Whisper backend: %s", e)
raise
# Load domain-specific initial prompt
self.initial_prompt = self._load_prompt()
# Load Silero VAD
self._vad_model: Optional[torch.nn.Module] = None
self._get_speech_timestamps = None
try:
from silero_vad import load_silero_vad, get_speech_timestamps
self._vad_model = load_silero_vad().to(self.device)
self._get_speech_timestamps = get_speech_timestamps
logger.info("Silero VAD loaded (threshold=%.2f).", VAD_THRESHOLD)
except Exception as exc:
logger.warning("Silero VAD could not be loaded: %s", exc)
def _load_prompt(self) -> str:
prompt_path = Path(__file__).parent.parent.parent / "prompts" / "whisper_initial_prompt.txt"
if prompt_path.exists():
prompt = prompt_path.read_text(encoding="utf-8").strip()
logger.info("Domain prompt loaded (%d characters).", len(prompt))
return prompt
return ""
def _load(self, audio_path: Union[str, Path]) -> torch.Tensor:
import soundfile as sf
audio_data, sr = sf.read(audio_path)
# Convert to torch tensor
waveform = torch.from_numpy(audio_data).float()
# Handle multi-channel (soundfile returns [samples, channels])
if len(waveform.shape) > 1:
waveform = torch.mean(waveform, dim=1)
# Resample if necessary
if sr != TARGET_SR:
import torchaudio.transforms as T
resampler = T.Resample(sr, TARGET_SR)
waveform = resampler(waveform.unsqueeze(0)).squeeze(0)
return waveform
def _apply_vad(self, audio: torch.Tensor) -> torch.Tensor:
if self._vad_model is None:
return audio
speech_timestamps = self._get_speech_timestamps(
audio, self._vad_model, sampling_rate=TARGET_SR, threshold=VAD_THRESHOLD
)
if not speech_timestamps:
return torch.tensor([], device=audio.device)
chunks = [audio[ts['start']:ts['end']] for ts in speech_timestamps]
return torch.cat(chunks) if chunks else torch.tensor([], device=audio.device)
def transcribe(self, audio_path: Union[str, Path]) -> str:
"""
Pure, high-performance transcription.
Returns a clean, single-stream Egyptian Arabic transcript.
"""
logger.info("Transcribing: %s", audio_path)
audio = self._load(audio_path).to(self.device)
audio_clean = self._apply_vad(audio)
if len(audio_clean) == 0:
logger.warning("No speech detected by VAD.")
return ""
generate_kwargs: dict = {"language": "arabic"}
if self.initial_prompt:
prompt_ids = self.processor.get_prompt_ids(self.initial_prompt)
if isinstance(prompt_ids, np.ndarray):
prompt_ids = torch.from_numpy(prompt_ids).to(self.device)
generate_kwargs["prompt_ids"] = prompt_ids
# Inference using the optimized pipeline
result = self.pipe(
audio_clean.cpu().numpy(),
generate_kwargs=generate_kwargs,
)
transcript = result.get("text", "").strip()
logger.info("Transcription complete. Length: %d chars.", len(transcript))
return transcript
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
# Quick test with merged model if available
t = WhisperTranscriber(model_path="outputs/checkpoints/merged_model")
print(t.transcribe("1775560189.41808.wav"))
|