Update: Refactored the code

src/transcriber.py

@@ -1,40 +1,63 @@
-google_colab_gpu = "cuda" if torch.cuda.is_available() else "cpu"
+import gc
+import json
 
-def transcribe_wav_file(wav_filename, index):
-    batch_size = 16
-    compute_type = "float16"
+import torch
+import whisperx
 
-    model = whisperx.load_model("large-v2", device=google_colab_gpu, compute_type=compute_type)
 
-    audio = whisperx.load_audio(wav_filename)
-    result = model.transcribe(audio, batch_size=batch_size)
+class Transcriber:
 
-    model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=google_colab_gpu)
-    result = whisperx.align(result["segments"], model_a, metadata, audio, google_colab_gpu, return_char_alignments=False)
+    def __init__(self, language_code="en", device="cuda", compute_type="float16"):
+        self.device = device if torch.cuda.is_available() else "cpu"
 
-    # clear GPU memory
-    del model
-    del model_a
-    del metadata
-    gc.collect()
-    torch.cuda.empty_cache()
+        # Initialize the transcription model
+        self.transcription_model = whisperx.load_model(
+            "large-v2", device=self.device, compute_type=compute_type)
 
-    diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', device=google_colab_gpu)
+        # Initialize the alignment model
+        self.alignment_model, self.metadata = whisperx.load_align_model(
+            language_code=language_code, device=self.device)
+        self.transcript_output = None
 
-    diarize_segments = diarize_model(audio)
-    result = whisperx.assign_word_speakers(diarize_segments, result)
+    def transcribe(self, wav_filename, batch_size=16):
+        # Load the audio file
+        audio = whisperx.load_audio(wav_filename)
 
-    # clear GPU memory
-    del diarize_model
-    del diarize_segments
-    gc.collect()
-    torch.cuda.empty_cache()
+        # Transcribe the audio file
+        result = self.transcription_model.transcribe(audio, batch_size=batch_size)
 
-    # Save full results to a json file
-    with open(f'result_{index+1}.json', 'w') as f:
-        json.dump(result, f)
+        # Use the alignment model
+        result = whisperx.align(
+            result["segments"], self.alignment_model, self.metadata, audio, self.device,
+            return_char_alignments=False)
 
-    # Save just the transcripts to a .txt file
-    transcripts = [seg['value'] for seg in result['segments']]
-    with open(f'transcripts_{index+1}.txt', 'w') as f:
-        f.write('\n'.join(transcripts))
+        # 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()