Update: Applied Black formatting
This commit is contained in:
+28
-13
@@ -14,7 +14,12 @@ logger = logging.getLogger(__name__)
|
||||
class Transcriber:
|
||||
"""A class to transcribe audio files using whisperx."""
|
||||
|
||||
def __init__(self, language_code: str = "en", device: str = "cuda", compute_type: str = "float16"):
|
||||
def __init__(
|
||||
self,
|
||||
language_code: str = "en",
|
||||
device: str = "cuda",
|
||||
compute_type: str = "float16",
|
||||
):
|
||||
"""
|
||||
Init method.
|
||||
:param language_code: str: Language code for whisperx transcription (default is "en").
|
||||
@@ -24,14 +29,17 @@ class Transcriber:
|
||||
self.device = device if torch.cuda.is_available() else "cpu"
|
||||
# Initialize the transcription and alignment models
|
||||
try:
|
||||
self.transcription_model = whisperx.load_model("large-v2", device=self.device, compute_type=compute_type)
|
||||
self.alignment_model, self.metadata = whisperx.load_align_model(language_code=language_code,
|
||||
device=self.device)
|
||||
self.transcription_model = whisperx.load_model(
|
||||
"large-v2", device=self.device, compute_type=compute_type
|
||||
)
|
||||
self.alignment_model, self.metadata = whisperx.load_align_model(
|
||||
language_code=language_code, device=self.device
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Failed to initialize models: %s", e)
|
||||
raise
|
||||
self.transcript_output = None
|
||||
self.hf_token = os.environ.get('HF_TOKEN')
|
||||
self.hf_token = os.environ.get("HF_TOKEN")
|
||||
|
||||
def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict:
|
||||
"""
|
||||
@@ -48,15 +56,22 @@ class Transcriber:
|
||||
result = self.transcription_model.transcribe(audio, batch_size=batch_size)
|
||||
|
||||
# Use the alignment model
|
||||
result = whisperx.align(result["segments"], self.alignment_model, self.metadata, audio, self.device,
|
||||
return_char_alignments=False)
|
||||
result = whisperx.align(
|
||||
result["segments"],
|
||||
self.alignment_model,
|
||||
self.metadata,
|
||||
audio,
|
||||
self.device,
|
||||
return_char_alignments=False,
|
||||
)
|
||||
|
||||
# Clear GPU memory
|
||||
self._clear_memory()
|
||||
|
||||
# Run Diarization
|
||||
diarize_model = whisperx.DiarizationPipeline(use_auth_token=self.hf_token,
|
||||
device=self.device)
|
||||
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)
|
||||
|
||||
@@ -76,13 +91,13 @@ class Transcriber:
|
||||
"""
|
||||
try:
|
||||
# Save full results to a json file
|
||||
with open(output_json_file, 'w') as f:
|
||||
with open(output_json_file, "w") as f:
|
||||
json.dump(self.transcript_output, f)
|
||||
|
||||
# Save just the transcripts to a .txt file
|
||||
transcripts = [seg['text'] for seg in self.transcript_output['segments']]
|
||||
with open(output_txt_file, 'w') as f:
|
||||
f.write('\n'.join(transcripts))
|
||||
transcripts = [seg["text"] for seg in self.transcript_output["segments"]]
|
||||
with open(output_txt_file, "w") as f:
|
||||
f.write("\n".join(transcripts))
|
||||
except Exception as e:
|
||||
logger.exception("Failed to save transcript: %s", e)
|
||||
raise
|
||||
|
||||
Reference in New Issue
Block a user