Refactoring: adapted for changes in cuda, updated README.md, improved logger and added sentence level transcript generation.

Signed-off-by: Orihara Izaya <xloki.stormbringery@gmail.com>

src/transcriber.py

@@ -1,42 +1,61 @@
 import gc
-import json
 import logging
 import os
 
 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."""
+    """
+    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,
     ):
-        """
-        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").
-        """
+        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"
-        # Initialize the transcription and alignment models
+        self.logger.info("Using device: %s", self.device)
+
         try:
             self.transcription_model = whisperx.load_model(
-                "large-v2", device=self.device, compute_type=compute_type
+                "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:
-            logger.exception("Failed to initialize models: %s", e)
+            self.logger.exception("Failed to initialize models: %s", e)
             raise
         self.transcript_output = None
         self.hf_token = os.environ.get("HF_TOKEN")
@@ -44,18 +63,24 @@ class Transcriber:
     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.
+
+        Parameters:
+        ----------
+        wav_filename : str
+            Path to the wav file.
+        batch_size : int (default: 16)
+            Batch size for whisperx transcription.
+
+        Returns:
+        -------
+        dict
+            Transcription 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)
 
-            # Use the alignment model
             result = whisperx.align(
                 result["segments"],
                 self.alignment_model,
@@ -65,45 +90,22 @@ class Transcriber:
                 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_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
         except Exception as e:
-            logger.exception("Failed to transcribe: %s", e)
+            self.logger.exception("Failed to transcribe: %s", e)
             raise
 
-    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["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
-
-    @staticmethod
-    def _clear_memory():
+    def _clear_memory(self):
         """
         Clear GPU memory.
         """
@@ -111,5 +113,5 @@ class Transcriber:
             gc.collect()
             torch.cuda.empty_cache()
         except Exception as e:
-            logger.exception("Failed to clear GPU memory: %s", e)
+            self.logger.exception("Failed to clear GPU memory: %s", e)
             raise