Update: Refactored the code

README.md

@@ -4,13 +4,8 @@
 
 Transcriber code
 
-Currently very unorganized
-
 ## to-do
 
-- clean the code
-- organize the stuff
-- add functions for other data sources
 - seperate data production form code
 
 ```

src/file_manager.py

@@ -1,41 +1,73 @@
+import logging
 import os
 
 from pydub import AudioSegment
-from pytube import YouTube
+from pytube import YouTube, exceptions as pytube_exceptions
+
+logging.basicConfig(level=logging.INFO)
 
 
 class FileManager:
+    """
+    FileManager class that handles operations such as downloading and converting audio.
+    """
+
     def __init__(self):
         pass
 
-    def download_youtube_audio(self, youtube_link, filename):
+    @staticmethod
-        yt = YouTube(youtube_link)
+    def download_youtube_audio(youtube_link: str, filename: str):
+        """
+        Downloads audio from a YouTube link and saves it in the location specified by filename.
+        """
+        try:
+            yt = YouTube(youtube_link)
+        except pytube_exceptions.PytubeError as e:
+            logging.error("An Error occurred while fetching the YouTube video: " + str(e))
+            return
 
-        # Get the highest quality audio stream
         audio_stream = yt.streams.get_audio_only()
 
-        # Download if the file doesn't already exist
         if not os.path.exists(filename):
-            audio_stream.download(filename=filename)
+            try:
+                audio_stream.download(filename=filename)
+                logging.info(f"Downloaded {youtube_link} successfully.")
+            except Exception as e:
+                logging.error("An Error occurred while downloading the YouTube video: " + str(e))
+                return
 
-    def convert_to_wav(self, mp4_file, wav_file):
+    @staticmethod
-        sound = AudioSegment.from_file(mp4_file, format="mp4")
+    def convert_to_wav(mp4_file: str, wav_file: str):
-        sound.export(wav_file, format="wav")
+        """
+        Converts a mp4 audio file into wav format.
+        """
+        try:
+            sound = AudioSegment.from_file(mp4_file, format="mp4")
+            sound.export(wav_file, format="wav")
+        except Exception as e:
+            logging.error(f"An error occurred while converting {mp4_file} to wav: {e}")
 
-    def load_local_file(self, local_file_path):
+    @staticmethod
-        # Check the file type
+    def load_local_file(local_file_path: str) -> str:
+        """
+        Loads a local file and converts it to wav format.
+        Return wav file path.
+        """
         ext = os.path.splitext(local_file_path)[-1].lower()
 
-        if ext == ".mp4":
+        if ext in [".mp4", ".mp3"]:
-            wav_file = local_file_path.replace(".mp4", ".wav")
+            wav_file = local_file_path.replace(ext, ".wav")
-            self.convert_to_wav(local_file_path, wav_file)
+            try:
-            return wav_file
+                audio = AudioSegment.from_file(local_file_path, format=ext[1:])
-        elif ext == ".mp3":
+                audio.export(wav_file, format="wav")
-            wav_file = local_file_path.replace(".mp3", ".wav")
+            except Exception as e:
-            audio = AudioSegment.from_file(local_file_path, format="mp3")
+                logging.error(f"An error occurred while converting {local_file_path} to wav: {e}")
-            audio.export(wav_file, format="wav")
+
             return wav_file
+
         elif ext == ".wav":
             return local_file_path
+
         else:
+            logging.error(f"Unsupported file format: {ext}.")
             raise ValueError("Unsupported file format.")

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 and alignment models
-        # Initialize the transcription model
+        try:
-        self.transcription_model = whisperx.load_model(
+            self.transcription_model = whisperx.load_model("large-v2", device=self.device, compute_type=compute_type)
-            "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)
-        # Initialize the alignment model
+        except Exception as e:
-        self.alignment_model, self.metadata = whisperx.load_align_model(
+            logger.exception("Failed to initialize models: %s", e)
-            language_code=language_code, device=self.device)
+            raise
         self.transcript_output = None
 
-    def transcribe(self, wav_filename, batch_size=16):
+    def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict:
-        # Load the audio file
+        """
-        audio = whisperx.load_audio(wav_filename)
+        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
+            # Transcribe the audio file
-        result = self.transcription_model.transcribe(audio, batch_size=batch_size)
+            result = self.transcription_model.transcribe(audio, batch_size=batch_size)
 
-        # Use the alignment model
+            # Use the alignment model
-        result = whisperx.align(
+            result = whisperx.align(result["segments"], self.alignment_model, self.metadata, audio, self.device,
-            result["segments"], self.alignment_model, self.metadata, audio, self.device,
+                                    return_char_alignments=False)
-            return_char_alignments=False)
 
-        # Clear GPU memory
+            # Clear GPU memory
-        self._clear_memory()
+            self._clear_memory()
 
-        # Run Diarization
+            # Run Diarization
-        diarize_model = whisperx.DiarizationPipeline(use_auth_token='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',
+            diarize_model = whisperx.DiarizationPipeline(use_auth_token='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',
-                                                     device=self.device)
+                                                         device=self.device)
-        diarize_segments = diarize_model(audio)
+            diarize_segments = diarize_model(audio)
-        result = whisperx.assign_word_speakers(diarize_segments, result)
+            result = whisperx.assign_word_speakers(diarize_segments, result)
 
-        # Clear GPU memory again
+            # Clear GPU memory again
-        self._clear_memory()
+            self._clear_memory()
-        self.transcript_output = result
+            self.transcript_output = result
-        return 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):
+    def save_transcript(self, output_json_file: str, output_txt_file: str):
-        # Save full results to a json file
+        """
-        with open(output_json_file, 'w') as f:
+        Save the transcription result to a json file and a text file.
-            json.dump(self.transcript_output, f)
+        :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
+            # Save just the transcripts to a .txt file
-        transcripts = [seg['value'] for seg in self.transcript_output['segments']]
+            transcripts = [seg['value'] for seg in self.transcript_output['segments']]
-        with open(output_txt_file, 'w') as f:
+            with open(output_txt_file, 'w') as f:
-            f.write('\n'.join(transcripts))
+                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
+        Clear GPU memory.
-        del self.metadata
+        """
-        gc.collect()
+        try:
-        torch.cuda.empty_cache()
+            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

src/transcript_generator.py

@@ -1,27 +1,46 @@
+import logging
+from typing import List
+
 from src.file_manager import FileManager
 from src.transcriber import Transcriber
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 class TranscriptGenerator:
+    """
+    TranscriptGenerator class that coordinates FileManager and Transcriber classes
+    to generate transcriptions from given input sources.
+    """
+
     def __init__(self):
         self.file_manager = FileManager()
         self.transcriber = Transcriber()
 
-    def generate(self, input_sources):
+    def generate(self, input_sources: List[str]) -> None:
+        """
+        Generates transcriptions for the given list of input sources.
+
+        Args:
+            input_sources (List[str]): List of input sources which could be URLs or local file paths.
+        """
         for i, source in enumerate(input_sources):
-            if source.startswith('http'):
+            try:
-                mp4_filename = f"audio_{i + 1}.mp4"
+                if source.startswith('http'):
-                wav_filename = f"audio_{i + 1}.wav"
+                    mp4_filename = f"audio_{i + 1}.mp4"
+                    wav_filename = f"audio_{i + 1}.wav"
+                    self.file_manager.download_youtube_audio(source, mp4_filename)
+                    self.file_manager.convert_to_wav(mp4_filename, wav_filename)
+                else:
+                    wav_filename = self.file_manager.load_local_file(source)
+                self.transcriber.transcribe(wav_filename)
+                json_output_file = f'result_{i + 1}.json'
+                txt_output_file = f'transcripts_{i + 1}.txt'
+                self.transcriber.save_transcript(json_output_file, txt_output_file)
+            except Exception as e:
+                logger.exception(f"Failed to generate transcript for source {source}. Error: {str(e)}")
+                continue
 
-                self.file_manager.download_youtube_audio(source, mp4_filename)
+        logger.info('Transcription generation completed.')
-                self.file_manager.convert_to_wav(mp4_filename, wav_filename)
-            else:
-                wav_filename = self.file_manager.load_local_file(source)
-
-            transcription = self.transcriber.transcribe(wav_filename)
-
-            json_output_file = f'result_{i + 1}.json'
-            txt_output_file = f'transcripts_{i + 1}.txt'
-            self.transcriber.save_transcript(json_output_file, txt_output_file)
-
-        print('Transcription generation completed.')