118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
import gc
|
|
import logging
|
|
import os
|
|
|
|
import torch
|
|
import whisperx
|
|
|
|
|
|
class Transcriber:
|
|
"""
|
|
A class to transcribe audio files using whisperx.
|
|
|
|
Attributes:
|
|
----------
|
|
language_code : str (default: "en")
|
|
Language code for whisperx transcription. Default is english.
|
|
device : str (default: "cuda")
|
|
Device to use for transcription. Default is "cuda" if available, else "cpu".
|
|
compute_type : str (default: "float16")
|
|
Precision type for whisperx transcription.
|
|
|
|
Methods
|
|
-------
|
|
transcribe(wav_filename: str, batch_size: int = 16) -> dict:
|
|
Transcribe the wav file.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
language_code: str = "en",
|
|
device: str = "cuda",
|
|
compute_type: str = "float16",
|
|
asr_options=None,
|
|
):
|
|
if asr_options is None:
|
|
asr_options = {
|
|
"repetition_penalty": 1,
|
|
"prompt_reset_on_temperature": 0.5,
|
|
"no_repeat_ngram_size": 0,
|
|
}
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
self.device = device if torch.cuda.is_available() else "cpu"
|
|
self.logger.info("Using device: %s", self.device)
|
|
|
|
try:
|
|
self.transcription_model = whisperx.load_model(
|
|
"large-v2",
|
|
device=self.device,
|
|
compute_type=compute_type,
|
|
language=language_code,
|
|
asr_options=asr_options,
|
|
)
|
|
self.alignment_model, self.metadata = whisperx.load_align_model(
|
|
language_code=language_code, device=self.device
|
|
)
|
|
except Exception as e:
|
|
self.logger.exception("Failed to initialize models: %s", e)
|
|
raise
|
|
self.transcript_output = None
|
|
self.hf_token = os.environ.get("HF_TOKEN")
|
|
|
|
def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict:
|
|
"""
|
|
Transcribe the wav file.
|
|
|
|
Parameters:
|
|
----------
|
|
wav_filename : str
|
|
Path to the wav file.
|
|
batch_size : int (default: 16)
|
|
Batch size for whisperx transcription.
|
|
|
|
Returns:
|
|
-------
|
|
dict
|
|
Transcription result.
|
|
"""
|
|
try:
|
|
audio = whisperx.load_audio(wav_filename)
|
|
|
|
result = self.transcription_model.transcribe(audio, batch_size=batch_size)
|
|
|
|
result = whisperx.align(
|
|
result["segments"],
|
|
self.alignment_model,
|
|
self.metadata,
|
|
audio,
|
|
self.device,
|
|
return_char_alignments=False,
|
|
)
|
|
|
|
self._clear_memory()
|
|
|
|
diarize_model = whisperx.DiarizationPipeline(
|
|
use_auth_token=self.hf_token, device=self.device
|
|
)
|
|
diarize_segments = diarize_model(audio)
|
|
result = whisperx.assign_word_speakers(diarize_segments, result)
|
|
|
|
self._clear_memory()
|
|
self.transcript_output = result
|
|
return result
|
|
except Exception as e:
|
|
self.logger.exception("Failed to transcribe: %s", e)
|
|
raise
|
|
|
|
def _clear_memory(self):
|
|
"""
|
|
Clear GPU memory.
|
|
"""
|
|
try:
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
self.logger.exception("Failed to clear GPU memory: %s", e)
|
|
raise
|