Update: Refactored the code

src/transcriber.py

@@ -1,63 +1,100 @@
 import gc
 import json
+import logging
 
 import torch
 import whisperx
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 class Transcriber:
+    """A class to transcribe audio files using whisperx."""
 
-    def __init__(self, language_code="en", device="cuda", compute_type="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").
+        :param device: str: Device to use for transcription (default is "cuda" if available, else "cpu").
+        :param compute_type: str: Compute type for whisperx transcription (default is "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)
+        # 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)
+        except Exception as e:
+            logger.exception("Failed to initialize models: %s", e)
+            raise
         self.transcript_output = None
 
-    def transcribe(self, wav_filename, batch_size=16):
-        # Load the audio file
-        audio = whisperx.load_audio(wav_filename)
+    def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict:
+        """
+        Transcribe the wav file.
+        :param wav_filename: str: Path to the wav file.
+        :param batch_size: int: Batch size for whisperx transcription (default is 16).
+        :return: dict: Transcribed and aligned result.
+        """
+        try:
+            # Load the audio file
+            audio = whisperx.load_audio(wav_filename)
 
-        # Transcribe the audio file
-        result = self.transcription_model.transcribe(audio, batch_size=batch_size)
+            # 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)
+            # 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()
+            # 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)
+            # 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
+            # Clear GPU memory again
+            self._clear_memory()
+            self.transcript_output = result
+            return result
+        except Exception as e:
+            logger.exception("Failed to transcribe: %s", e)
+            raise
 
-    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)
+    def save_transcript(self, output_json_file: str, output_txt_file: str):
+        """
+        Save the transcription result to a json file and a text file.
+        :param output_json_file: str: Path to the output json file.
+        :param output_txt_file: str: Path to the output text file.
+        """
+        try:
+            # 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))
+            # 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))
+        except Exception as e:
+            logger.exception("Failed to save transcript: %s", e)
+            raise
 
     def _clear_memory(self):
-        del self.transcription_model
-        del self.alignment_model
-        del self.metadata
-        gc.collect()
-        torch.cuda.empty_cache()
+        """
+        Clear GPU memory.
+        """
+        try:
+            del self.transcription_model
+            del self.alignment_model
+            del self.metadata
+            gc.collect()
+            torch.cuda.empty_cache()
+        except Exception as e:
+            logger.exception("Failed to clear GPU memory: %s", e)
+            raise