Update: Refactored the code

2023-10-06 15:45:33 +02:00
parent 0b1c24de2a
commit 57cc73a439
4 changed files with 165 additions and 82 deletions
@@ -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):
-        yt = YouTube(youtube_link)
+    @staticmethod
+    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):
-        sound = AudioSegment.from_file(mp4_file, format="mp4")
-        sound.export(wav_file, format="wav")
+    @staticmethod
+    def convert_to_wav(mp4_file: str, wav_file: str):
+        """
+        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):
-        # Check the file type
+    @staticmethod
+    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":
-            wav_file = local_file_path.replace(".mp4", ".wav")
-            self.convert_to_wav(local_file_path, wav_file)
-            return wav_file
-        elif ext == ".mp3":
-            wav_file = local_file_path.replace(".mp3", ".wav")
-            audio = AudioSegment.from_file(local_file_path, format="mp3")
-            audio.export(wav_file, format="wav")
+        if ext in [".mp4", ".mp3"]:
+            wav_file = local_file_path.replace(ext, ".wav")
+            try:
+                audio = AudioSegment.from_file(local_file_path, format=ext[1:])
+                audio.export(wav_file, format="wav")
+            except Exception as e:
+                logging.error(f"An error occurred while converting {local_file_path} to wav: {e}")
+
            return wav_file
+
        elif ext == ".wav":
            return local_file_path
+
        else:
+            logging.error(f"Unsupported file format: {ext}.")
            raise ValueError("Unsupported file format.")
@@ -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
@@ -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'):
-                mp4_filename = f"audio_{i + 1}.mp4"
-                wav_filename = f"audio_{i + 1}.wav"
+            try:
+                if source.startswith('http'):
+                    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)
-                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.')
+        logger.info('Transcription generation completed.')