64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import gc
|
|
import json
|
|
|
|
import torch
|
|
import whisperx
|
|
|
|
|
|
class Transcriber:
|
|
|
|
def __init__(self, language_code="en", device="cuda", compute_type="float16"):
|
|
self.device = device if torch.cuda.is_available() else "cpu"
|
|
|
|
# Initialize the transcription model
|
|
self.transcription_model = whisperx.load_model(
|
|
"large-v2", device=self.device, compute_type=compute_type)
|
|
|
|
# Initialize the alignment model
|
|
self.alignment_model, self.metadata = whisperx.load_align_model(
|
|
language_code=language_code, device=self.device)
|
|
self.transcript_output = None
|
|
|
|
def transcribe(self, wav_filename, batch_size=16):
|
|
# Load the audio file
|
|
audio = whisperx.load_audio(wav_filename)
|
|
|
|
# Transcribe the audio file
|
|
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)
|
|
|
|
# Clear GPU memory
|
|
self._clear_memory()
|
|
|
|
# Run Diarization
|
|
diarize_model = whisperx.DiarizationPipeline(use_auth_token='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',
|
|
device=self.device)
|
|
diarize_segments = diarize_model(audio)
|
|
result = whisperx.assign_word_speakers(diarize_segments, result)
|
|
|
|
# Clear GPU memory again
|
|
self._clear_memory()
|
|
self.transcript_output = result
|
|
return result
|
|
|
|
def save_transcript(self, output_json_file, output_txt_file):
|
|
# Save full results to a json file
|
|
with open(output_json_file, 'w') as f:
|
|
json.dump(self.transcript_output, f)
|
|
|
|
# Save just the transcripts to a .txt file
|
|
transcripts = [seg['value'] for seg in self.transcript_output['segments']]
|
|
with open(output_txt_file, 'w') as f:
|
|
f.write('\n'.join(transcripts))
|
|
|
|
def _clear_memory(self):
|
|
del self.transcription_model
|
|
del self.alignment_model
|
|
del self.metadata
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|