107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
import threading
|
|
import whisper
|
|
import asyncio
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import librosa
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Dict
|
|
|
|
from logging import getLogger
|
|
from whisper import Whisper
|
|
|
|
from ielts_be.services import ISpeechToTextService
|
|
|
|
"""
|
|
The whisper model is not thread safe, a thread pool
|
|
with 4 whisper models will be created so it can
|
|
process up to 4 transcriptions at a time.
|
|
|
|
The base model requires ~1GB so 4 instances is the safe bet:
|
|
https://github.com/openai/whisper?tab=readme-ov-file#available-models-and-languages
|
|
"""
|
|
class OpenAIWhisper(ISpeechToTextService):
|
|
def __init__(self, model_name: str = "base", num_models: int = 4):
|
|
self._model_name = model_name
|
|
self._num_models = num_models
|
|
self._models: Dict[int, 'Whisper'] = {}
|
|
self._lock = threading.Lock()
|
|
self._next_model_id = 0
|
|
self._is_closed = False
|
|
self._logger = getLogger(__name__)
|
|
|
|
for i in range(num_models):
|
|
self._models[i] = whisper.load_model(self._model_name, in_memory=True)
|
|
|
|
self._executor = ThreadPoolExecutor(
|
|
max_workers=num_models,
|
|
thread_name_prefix="whisper_worker"
|
|
)
|
|
|
|
def get_model(self) -> 'Whisper':
|
|
with self._lock:
|
|
model_id = self._next_model_id
|
|
self._next_model_id = (self._next_model_id + 1) % self._num_models
|
|
return self._models[model_id]
|
|
|
|
async def speech_to_text(self, path: str) -> str:
|
|
def transcribe():
|
|
try:
|
|
audio, sr = sf.read(path)
|
|
|
|
# Convert to mono first to reduce memory usage
|
|
if len(audio.shape) > 1:
|
|
audio = audio.mean(axis=1)
|
|
|
|
# Resample from 48kHz to 16kHz
|
|
audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
|
|
|
|
# Normalize to [-1, 1] range
|
|
audio = audio.astype(np.float32)
|
|
if np.max(np.abs(audio)) > 0:
|
|
audio = audio / np.max(np.abs(audio))
|
|
|
|
# Break up long audio into chunks (30 seconds at 16kHz = 480000 samples)
|
|
max_samples = 480000
|
|
if len(audio) > max_samples:
|
|
chunks = []
|
|
for i in range(0, len(audio), max_samples):
|
|
chunk = audio[i:i + max_samples]
|
|
chunks.append(chunk)
|
|
|
|
model = self.get_model()
|
|
texts = []
|
|
for chunk in chunks:
|
|
result = model.transcribe(
|
|
chunk,
|
|
fp16=False,
|
|
language='English',
|
|
verbose=False
|
|
)["text"]
|
|
texts.append(result)
|
|
return " ".join(texts)
|
|
else:
|
|
model = self.get_model()
|
|
return model.transcribe(
|
|
audio,
|
|
fp16=False,
|
|
language='English',
|
|
verbose=False
|
|
)["text"]
|
|
|
|
except Exception as e:
|
|
raise
|
|
|
|
loop = asyncio.get_running_loop()
|
|
return await loop.run_in_executor(self._executor, transcribe)
|
|
|
|
def close(self):
|
|
with self._lock:
|
|
if not self._is_closed:
|
|
self._is_closed = True
|
|
if self._executor:
|
|
self._executor.shutdown(wait=True, cancel_futures=True)
|
|
|
|
def __del__(self):
|
|
self.close()
|