Update: Refactored the code

.gitignore

@@ -0,0 +1,13 @@
+# PyCharm IDE
+.idea/
+*.iml
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/

2023-09-20_transcripts_Video-SRS.py

@@ -1,14 +0,0 @@
-from src.data_handling import download_audio, convert_to_wav
-from src.transcriber import transcribe_wav_file
-
-youtube_links = ['https://www.youtube.com/watch?v=BVVo1BaL6JA', 'https://www.youtube.com/watch?v=m0gr3snccVY&list=PL8p82vzFeTAj3cbHaXE9sC48pm4oZY2rH&index=2', 'https://www.youtube.com/watch?v=ThjgUfck8Sc&list=PL8p82vzFeTAj3cbHaXE9sC48pm4oZY2rH&index=3']
-#youtube_links = ['https://www.youtube.com/watch?v=6A5ZSiosgAk']
-
-for i, link in enumerate(youtube_links):
-    mp4_filename = f"audio_{i+1}.mp4"
-    wav_filename = f"audio_{i+1}.wav"
-
-    download_audio(link, mp4_filename)
-    convert_to_wav(mp4_filename, wav_filename)
-
-    transcribe_wav_file(wav_filename, i)

README.md

@@ -1,7 +1,5 @@
 # Project_whisperTrans01_Code
 
-
-
 ## Content
 
 Transcriber code
@@ -13,4 +11,34 @@ Currently very unorganized
 - clean the code
 - organize the stuff
 - add functions for other data sources
-- seperate data production form code
+- seperate data production form code
+
+```
+project_root/
+├── src/
+│   ├── file_manager.py   # Contains FileManager class
+│   ├── transcriber.py    # Contains Transcriber class
+│   ├── transcript_generator.py    # Contains TranscriptGenerator class
+│   ├── data_handling.py  # Contains download_audio and convert_to_wav methods
+│   └── main.py    # Main script that runs the process
+├── tests/    
+│   ├── test_file_manager.py
+│   ├── test_transcriber.py
+│   └── test_transcript_generator.py
+├── data/
+│   ├── audio/   # downloaded mp4 and output wav files
+│   └── transcripts/  # JSON and text transcript files
+└── README.md
+```
+
+- **src/**: This directory will contain all your source code. Inside the src folder, I added:
+    - **main.py**: This will be the main entry point for the program. It will use the `TranscriptGenerator` class which
+      orchestrates the use of `FileManager` and `Transcriber`.
+    - **file_manager.py**: Contains the `FileManager` class implementation.
+    - **transcriber.py**: Contains the `Transcriber` class implementation.
+    - **transcript_generator.py**: Contains the `TranscriptGenerator` class implementation.
+
+- **tests/**: This directory contains all your test scripts corresponding to your source files.
+
+- **data/**: This directory contains all the data needed for your code. For simplicity, I categorize the files into
+  separate directories for audio and transcript files.

src/data_handling.py

@@ -1,13 +0,0 @@
-def download_audio(youtube_link, filename):
-    yt = YouTube(youtube_link)
-
-    # get highest quality audio stream
-    audio_stream = yt.streams.get_audio_only()
-
-    # download
-    if not os.path.exists(filename):
-        audio_stream.download(filename=filename)
-
-def convert_to_wav(mp4_file, wav_file):
-    sound = AudioSegment.from_file(mp4_file, format="mp4")
-    sound.export(wav_file, format="wav")

src/file_manager.py

@@ -0,0 +1,41 @@
+import os
+
+from pydub import AudioSegment
+from pytube import YouTube
+
+
+class FileManager:
+    def __init__(self):
+        pass
+
+    def download_youtube_audio(self, youtube_link, filename):
+        yt = YouTube(youtube_link)
+
+        # 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)
+
+    def convert_to_wav(self, mp4_file, wav_file):
+        sound = AudioSegment.from_file(mp4_file, format="mp4")
+        sound.export(wav_file, format="wav")
+
+    def load_local_file(self, local_file_path):
+        # Check the file type
+        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")
+            return wav_file
+        elif ext == ".wav":
+            return local_file_path
+        else:
+            raise ValueError("Unsupported file format.")

src/main.py

@@ -0,0 +1,16 @@
+from src.transcript_generator import TranscriptGenerator
+
+
+def main():
+    # input could be a YouTube link or a local file
+    input_sources = [
+        'https://www.youtube.com/watch?v=BVVo1BaL6JA',
+        '/path/to/your/local/file.mp3'
+    ]
+
+    transcript_generator = TranscriptGenerator()
+    transcript_generator.generate(input_sources)
+
+
+if __name__ == "__main__":
+    main()

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()

src/transcript_generator.py

@@ -0,0 +1,27 @@
+from src.file_manager import FileManager
+from src.transcriber import Transcriber
+
+
+class TranscriptGenerator:
+    def __init__(self):
+        self.file_manager = FileManager()
+        self.transcriber = Transcriber()
+
+    def generate(self, input_sources):
+        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"
+
+                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.')

tests/test_file_manager.py

@@ -0,0 +1,33 @@
+import os
+import unittest
+
+from src.file_manager import FileManager
+
+
+class FileManagerTest(unittest.TestCase):
+    def setUp(self):
+        self.file_manager = FileManager()
+        self.dir = '<Your file directory>'  # replace with your directory
+
+    def test_download_youtube_audio(self):
+        youtube_link = '<Your YouTube link>'  # replace with a YouTube link
+        filename = f'{self.dir}/audio.mp4'
+        self.file_manager.download_youtube_audio(youtube_link, filename)
+        self.assertTrue(os.path.exists(filename))
+        os.remove(filename)  # clean up after test
+
+    def test_convert_to_wav(self):
+        mp4_file = '<Your mp4 filename>'  # replace with an MP4 file in your directory
+        wav_file = mp4_file.replace('.mp4', '.wav')
+        self.file_manager.convert_to_wav(mp4_file, wav_file)
+        self.assertTrue(os.path.exists(wav_file))
+        os.remove(wav_file)  # clean up after test
+
+    def test_load_local_file(self):
+        local_file_path = '<Your mp4 filename>'  # replace with a file in your directory
+        result_path = self.file_manager.load_local_file(local_file_path)
+        self.assertEqual(result_path, local_file_path.replace('.mp4', '.wav'))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_transcriber.py

@@ -0,0 +1,32 @@
+import os
+import unittest
+
+from src.transcriber import Transcriber
+
+
+class TranscriberTest(unittest.TestCase):
+    def setUp(self):
+        self.transcriber = Transcriber()
+
+    def test_transcribe(self):
+        wav_file = '<Your wav file>'  # Replace with path to a .wav file on your machine
+        result = self.transcriber.transcribe(wav_file)
+        self.assertIsNotNone(result)
+
+    def test_save_transcript(self):
+        wav_file = '<Your wav file>'  # Replace with path to a .wav file on your machine
+        result = self.transcriber.transcribe(wav_file)
+
+        json_output_file = '<Path to your temporary JSON output file>'  # Temporary output file for testing
+        txt_output_file = '<Path to your temporary TXT output file>'  # Temporary output file for testing
+        self.transcriber.save_transcript(json_output_file, txt_output_file)
+
+        self.assertTrue(os.path.exists(json_output_file))
+        self.assertTrue(os.path.exists(txt_output_file))
+
+        os.remove(json_output_file)  # clean up after test
+        os.remove(txt_output_file)  # clean up after test
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_transcript_generator.py

@@ -0,0 +1,19 @@
+import unittest
+
+from src.transcript_generator import TranscriptGenerator
+
+
+class TranscriptGeneratorTest(unittest.TestCase):
+    def setUp(self):
+        self.transcript_generator = TranscriptGenerator()
+
+    def test_generate(self):
+        input_sources = [
+            'https://www.youtube.com/watch?v=BVVo1BaL6JA',
+            '/path/to/your/local/file.mp3'
+        ]
+        self.transcript_generator.generate(input_sources)
+
+
+if __name__ == "__main__":
+    unittest.main()