Documentation: Updated README.md with information on the strucutre of a job

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

.gitignore

@@ -11,6 +11,3 @@ __pycache__/
 env/
 venv/
 ENV/
-
-# Large files
-data/models # local storage of transcriber models

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Gerrit Anders
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,15 +1,27 @@
-# Project_whisperTrans01_Code
+# Project_whisperTrans_Code
 
-## Content
+## Overview
 
-Transcriber code
+whisperTrans is a video2transcript tool built upon [Whisper](https://github.com/openai/whisper)
+and [WhisperX](https://github.com/m-bain/whisperX). For automated speech recognition (ASR).
+The project provides an efficient implementation for running transcriptions locally on open source models from local
+files and YouTube videos. It is structured to handle large-scale audio transcription tasks with accuracy and ease using
+diarization and alignment techniques.
 
 ## Setup
 
-Before running this script, make sure to install the necessary python packages using the following command:
+The project is tested under python 3.10.12.
+In order to edit video material ffmpeg is required. Furthermore, machine learning models are employed in this project to
+convert audio content to text. Which models are used depends on the language of the video content. The models are
+downloaded on demand from the model hosting platform [huggingface](https://huggingface.co/). To do so you need to create
+an account and set the environment variable HF_TOKEN to your token. This is necessary to accept the conditions of the
+models used.
+
+Before running this script, make sure to install the necessary python packages. We recommend using a virtual environment
+to manage dependencies. You can install the required packages with the following command:
 
 ```bash
-pip install -r requirements.txt
+pip install -r requirements.txt  --index-url https://download.pytorch.org/whl/cu118
 ```
 
 If your system does not have ffmpeg installed you can install it via:
@@ -18,9 +30,44 @@ If your system does not have ffmpeg installed you can install it via:
 sudo apt install ffmpeg
 ```
 
-## Running the script
+For GPU usage you need to install the CUDA toolkit. You can install it via:
 
-To generate transcripts create a new job with your specifications (see jobs section below).
+```bash
+sudo apt install nvidia-cuda-toolkit-11-8
+```
+
+This might require to add the package the following way:
+
+```bash
+wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
+sudo dpkg -i cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
+sudo cp /var/cuda-repo-ubuntu2204-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
+```
+
+Afterward, set the environmental variables:
+
+```bash
+echo 'export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}' >> ~/.bashrc
+echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
+```
+
+and reload your `~/.bashrc` file to apply these changes:
+
+```bash
+source ~/.bashrc
+```
+
+## Prerequisites
+
+- Python 3.10
+- ffmpeg
+- Access to Hugging Face for model downloads
+- CUDA 11.8, and cuBLAS 11.x and cuDNN 8.x
+- CUDA-compatible GPU with <8GB VRAM (recommended for faster processing)
+
+## Usage
+
+To generate transcripts create a new job with your specifications (see [Jobs secion](#jobs) below).
 <br/>
 If you need to download the transcriber model set your huggingface token:
 
@@ -36,65 +83,72 @@ On Windows (Command prompt):
 set HF_TOKEN="your-token-here"
 ```
 
-Then run the following command to generate your transcript:
+After setting up the token, execute the script with:
 
 ```bash
-main.py --job {name}
+python main.py --job {name}
 ```
 
-where name is the name of your job.
+Here, `name` is the identifier for your transcription job.
+Note that you might need to go to huggingface to accept the conditions of the models used. Especially when running the
+transcriber for the first time on your huggingface account.
 
 ## Jobs
 
 Jobs specify tasks for transcript creation. A job contains the following information.
 
 - `name`: string: The name of the job
-- `source`: List[string]: The sources of the videos (paths, YouTube links, TibAV links)
-- `modelpath`: Optional[string]: The path to the model if already saved locally (will be downloaded temporarily if not)
-- `language`: Optional[string]: The language of the video (If not specified employs auto-detection)
+- `source`: List[string]: The sources of the videos (paths, YouTube links)
+- `language`: Optional[string]: The language of the video (If not specified employs auto-detection) as a language code
+  (e.g. 'en', 'de', 'fr')
+- `output_format`: dict: The requested output formats
+    - `get_json_version`: bool: Whether to get the full data json version of the transcript (including word level time
+      stamps)
+    - `get_txt_version`: bool: Whether to get a plain text version of the transcript
+    - `get_sentence_level_json_version`: bool: Whether to get a sentence level json version of the transcript
 
-## to-do
+## Contributing
 
-- seperate data production form code
-- remember to accept the conditions on diarization if you use huggingface
-- do a config json
-- do option to have models already stored and point to them (is this necessary? > if you do not want to interact with
-  huggingface yes)
-    - also for single language models
-- check warnings
-- downloader for tib-av
-- more optional setting
-- add progress bar
+Contributions are welcome! Please fork the repository and make your changes in a separate branch. Submit a pull request
+with your proposed modifications. For major changes, please open an issue first to discuss what you would like to
+change.
+
+## License
+
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
+
+## Contact
+
+This project is maintained by [Gerrit Anders](https://www.iwm-tuebingen.de/www/de/personen/ma.html?uid=ganders).
+For any questions or support, please open an issue on this repository or reach out via email
+at [g.anders@iwm-tuebingen.de](mailto:g.anders@iwm-tuebingen.de).
+
+## Future plans
+
+- **Tib-AV Downloader**: Add a downloader capability for [TibAV](https://av.tib.eu/) video links.
+- **Progress Bar**: Implement a progress bar for better user feedback during processing.
+- **Local Model Directory**: Allow models to be stored locally for offline use and faster processing.
+- **Extend Mock Tests**: Add actual references for tests.
 
 ## Structure
 
 ```
 project_root/
 ├── src/
-│   ├── file_manager.py   # Contains FileManager class
+│   ├── video_manager.py   # Contains VideoManager 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
+│   └── data_handling.py  # Contains download_audio and convert_to_wav methods
 ├── tests/    
 │   ├── test_file_manager.py
 │   ├── test_transcriber.py
 │   └── test_transcript_generator.py
 ├── data/
-│   ├── audio/   # downloaded mp4 and output wav files
-│   ├── models/   # models used for transcripts (are ignored by git and not necessarily)
-│   └── transcripts/  # JSON and text transcript files
+│   ├── logs/  # logs of previous jobs
+│   ├── jobs/   # place to store job files
+│   └── results/   # place to store results
+├── main.py    # Main script that runs the process
+├── requirements.txt
+├── LICENSE
 └── 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, model and transcript files.

data/jobs/example_job.json

@@ -0,0 +1,11 @@
+{
+  "name": "myjob1",
+  "source": [
+    "https://www.youtube.com/watch?v=PrGl3w2sAiw"
+  ],
+  "transcript_version": {
+    "get_json_version": true,
+    "get_txt_version": true,
+    "get_sentence_level_json_version": true
+  }
+}

main.py

@@ -1,19 +1,49 @@
+import json
+import logging
 import os
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
 
+from src.logger import setup_logger
 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"]
+    parser: ArgumentParser = ArgumentParser(
+        description="Transcript Generator for Videos"
+    )
+
+    parser.add_argument("--job", type=str, help="The name of the job", required=True)
+
+    arguments: Namespace = parser.parse_args()
+
+    with open(f"data/jobs/{arguments.job}", "r") as f:
+        job_config = json.load(f)
+
+    log_directory: Path = Path("data/logs")
+    os.makedirs(log_directory, exist_ok=True)
+    log_file_path: Path = log_directory / f"{arguments.job}.log"
+    setup_logger(__name__, log_file=log_file_path)
+    logger: logging.Logger = logging.getLogger(__name__)
+
+    logger.info("Starting the transcript generation.")
 
     hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
         hf_token = input("Please enter your HF_TOKEN: ")
     os.environ["HF_TOKEN"] = hf_token
 
-    transcript_generator = TranscriptGenerator()
-    transcript_generator.generate(input_sources)
+    transcript_generator: TranscriptGenerator = TranscriptGenerator(
+        job_name=job_config["name"],
+        language_code=job_config.get("language_code", "en"),
+        get_txt_version=job_config["transcript_version"]["get_txt_version"],
+        get_json_version=job_config["transcript_version"]["get_json_version"],
+        get_sentence_level_json=job_config["transcript_version"][
+            "get_sentence_level_json_version"
+        ],
+    )
+
+    transcript_generator.generate(job_config["source"])
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,7 +1,9 @@
-youtube-dl
-pytube
-pydub
-ffmpeg-python
-ffprobe
-torch
-git+https://github.com/m-bain/whisperx.git
+youtube-dl==2021.12.17
+pytube==15.0.0
+pydub==0.25.1
+ffmpeg-python==0.2.0
+ffprobe==0.5
+torch==2.0.0+cu118
+torchaudio==2.0.1+cu118
+pytorch-cuda=11.8
+git+https://github.com/m-bain/whisperx.git@d8a2b4f

src/file_manager.py

@@ -1,79 +0,0 @@
-import logging
-import os
-
-from pydub import AudioSegment
-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
-
-    @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
-
-        audio_stream = yt.streams.get_audio_only()
-
-        if not os.path.exists(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
-
-    @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}")
-
-    @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 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.")

src/logger.py

@@ -0,0 +1,47 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+
+def setup_logger(
+    name: str, level: int = logging.INFO, log_file: Optional[Path] = None
+) -> logging.Logger:
+    """
+    Configure and return a logger with the specified name and log level.
+
+    Parameters
+    ----------
+    name : str
+        The name of the logger.
+    level : int (default: logging.INFO)
+        The log level for the logger.
+    log_file : Path (default: None)
+        The path to the log file.
+    """
+    logger: logging.Logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    console_handler: logging.StreamHandler = logging.StreamHandler()
+    console_handler.setLevel(level)
+
+    if log_file:
+        file_handler: logging.FileHandler = logging.FileHandler(log_file)
+        file_handler.setLevel(level)
+    else:
+        file_handler = None
+
+    formatter: logging.Formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+
+    console_handler.setFormatter(formatter)
+
+    if file_handler:
+        file_handler.setFormatter(formatter)
+
+    logger.addHandler(console_handler)
+
+    if file_handler:
+        logger.addHandler(file_handler)
+
+    return logger

src/output_manager.py

@@ -0,0 +1,128 @@
+import json
+import logging
+
+
+class OutputManager:
+    """
+    A class to save the transcription result to a json file, a text file, and/or a sentence-level json file.
+
+    Attributes:
+    ----------
+        get_txt_version : bool
+            A flag to determine if the transcript should be saved in a text file.
+        get_json_version : bool
+            A flag to determine if the transcript should be saved in a JSON file.
+        get_sentence_level_json : bool
+            A flag to determine if the sentence-level transcript should be saved in a JSON file.
+
+    Methods
+    -------
+        save_transcript(transcript_output: dict, job_name: str, source_number: int) -> None:
+            Save the transcription result to a json file a text file and/or a sentence level json file.
+    """
+
+    def __init__(
+        self,
+        get_txt_version: bool,
+        get_json_version: bool,
+        get_sentence_level_json: bool,
+    ):
+        self.get_txt_version = get_txt_version
+        self.get_json_version = get_json_version
+        self.get_sentence_level_json = get_sentence_level_json
+
+        self.logger = logging.getLogger(__name__)
+
+    def save_transcript(
+        self, transcript_output: dict, job_name: str, source_number: int
+    ):
+        """
+        Save the transcription result to a json file a text file and/or a sentence level json file.
+
+        Parameters
+        ----------
+        transcript_output : dict
+            The transcription result.
+        job_name : str
+            The name of the job.
+        source_number : int
+            The enumeration of the source in the job.
+        """
+
+        json_output_file_name = f"result_{job_name}_{source_number + 1}.json"
+        txt_output_file_name = f"transcripts_{job_name}_{source_number + 1}.txt"
+        sentence_level_output_file_name = (
+            f"sentence_level_{job_name}_{source_number + 1}.json"
+        )
+        if self.get_json_version:
+            self._save_transcript_full(json_output_file_name, transcript_output)
+        if self.get_txt_version:
+            self._save_transcript_txt(txt_output_file_name, transcript_output)
+        if self.get_sentence_level_json:
+            self._save_transcript_sentence_level(
+                sentence_level_output_file_name, transcript_output
+            )
+        else:
+            self.logger.error("No output version specified. Not saving any output.")
+
+    def _save_transcript_full(self, output_json_file: str, transcript_output: dict):
+        """
+        Save the full transcription result to a json file.
+
+        Parameters
+        ----------
+        output_json_file : str
+            Path to the output json file.
+        transcript_output : dict
+            The transcription result to be saved.
+        """
+        try:
+            with open(output_json_file, "w") as f:
+                json.dump(transcript_output, f)
+        except Exception as e:
+            self.logger.exception("Failed to save transcript to json: %s", e)
+
+    def _save_transcript_txt(self, output_txt_file: str, transcript_output: dict):
+        """
+        Save the transcript to a text file.
+
+        Parameters
+        ----------
+        output_txt_file : str
+            Path to the output txt file.
+        transcript_output : dict
+            The transcription result to be saved.
+        """
+        try:
+            transcripts = [seg["text"] for seg in transcript_output["segments"]]
+            with open(output_txt_file, "w") as f:
+                f.write("\n".join(transcripts))
+        except Exception as e:
+            self.logger.exception("Failed to save transcript to txt: %s", e)
+
+    def _save_transcript_sentence_level(
+        self, output_json_file: str, transcript_output: dict
+    ):
+        """
+        Save the transcript to a sentence-level json file.
+
+        Parameters
+        ----------
+        output_json_file : str
+            Path to the output json file.
+        transcript_output : dict
+            The transcription result to be saved.
+        """
+        try:
+            sentence_level_transcript = []
+            for segment in transcript_output.get("segments", []):
+                sentence_info = {
+                    "start": segment.get("start"),
+                    "end": segment.get("end"),
+                    "text": segment.get("text"),
+                }
+                sentence_level_transcript.append(sentence_info)
+            with open(output_json_file, "w") as f:
+                json.dump(sentence_level_transcript, f, ensure_ascii=False, indent=4)
+        except Exception as e:
+            self.logger.exception("Failed to save sentence-level transcript: %s", e)

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

src/transcript_generator.py

@@ -1,48 +1,82 @@
 import logging
-from typing import List
 
-from src.file_manager import FileManager
+from src.output_manager import OutputManager
 from src.transcriber import Transcriber
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from src.video_manager import VideoManager
 
 
 class TranscriptGenerator:
     """
     TranscriptGenerator class that coordinates FileManager and Transcriber classes
-    to generate transcriptions from given input sources.
+    to generate transcriptions in desired formats from given input sources.
+
+    Attributes:
+    ----------
+        file_manager: FileManager
+            An instance of the FileManager class.
+        transcriber: Transcriber
+        job_name: str
+            The name of the job.
+        get_txt_version: bool (default: True)
+            A flag to determine if the transcript should be saved in a text file.
+        get_json_version: bool (default: True)
+            A flag to determine if the transcript should be saved in a JSON file.
+        get_sentence_level_json: bool (default: True)
+            A flag to determine if the sentence-level transcript should be saved in a JSON file.
+
+    Methods
+    -------
+        generate(input_sources: list[str]) -> None:
+            Generates transcriptions for the given list of input sources.
     """
 
-    def __init__(self):
-        self.file_manager = FileManager()
-        self.transcriber = Transcriber()
+    def __init__(
+        self,
+        job_name: str,
+        language_code: str = "en",
+        get_txt_version: bool = True,
+        get_json_version: bool = True,
+        get_sentence_level_json: bool = True,
+    ):
+        self.video_manager: VideoManager = VideoManager()
+        self.output_manager: OutputManager = OutputManager(
+            get_txt_version, get_json_version, get_sentence_level_json
+        )
+        self.transcriber: Transcriber = Transcriber(language_code)
+        self.job_name = job_name
 
-    def generate(self, input_sources: List[str]) -> None:
+        self.logger = logging.getLogger(__name__)
+
+    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.
+
+        Parameters
+        ----------
+            input_sources : list[str]
+                A list of input sources (local file paths or YouTube URLs) for which transcriptions should be generated.
         """
-        for i, source in enumerate(input_sources):
+        for source_number, source in enumerate(input_sources):
             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)
+                    mp4_filename = f"audio_{self.job_name}_{source_number + 1}.mp4"
+                    wav_filename = f"audio_{self.job_name}_{source_number + 1}.wav"
+                    self.video_manager.download_youtube_audio(source, mp4_filename)
+                    self.video_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)
+                    wav_filename = self.video_manager.load_local_file(source)
+                transcript_data = self.transcriber.transcribe(wav_filename)
+                self.output_manager.save_transcript(
+                    transcript_data, self.job_name, source_number
+                )
+
+                self.logger.info(f"Transcription generated for source {source}.")
+
             except Exception as e:
-                logger.exception(
+                self.logger.exception(
                     f"Failed to generate transcript for source {source}. Error: {str(e)}"
                 )
                 continue
 
-        logger.info("Transcription generation completed.")
+        self.logger.info("Transcription generation completed.")

src/video_manager.py

@@ -0,0 +1,111 @@
+import logging
+import os
+
+from pydub import AudioSegment
+from pytube import YouTube, exceptions as pytube_exceptions
+
+
+class VideoManager:
+    """
+    FileManager class that handles operations on input videos such as downloading and converting to audio.
+    To arrive at the correct input format for transcript generation.
+
+    Attributes:
+    ----------
+        None
+
+    Methods
+    -------
+        download_youtube_audio(youtube_link: str, filename: str) -> None:
+            Downloads audio from a YouTube link and saves it in the location specified by filename.
+        convert_to_wav(mp4_file: str, wav_file: str) -> None:
+            Converts a mp4 audio file into wav format.
+        load_local_file(local_file_path: str) -> str:
+            Loads a local file and converts it to wav format
+    """
+
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+
+    def download_youtube_audio(self, youtube_video_link: str, filename: str) -> None:
+        """
+        Downloads audio from a YouTube link and saves it in the location specified by filename.
+
+        Parameters
+        ----------
+            youtube_video_link : str
+                The YouTube video link from which the audio should be downloaded.
+            filename : str
+                The location where the audio should be saved
+        """
+        try:
+            yt = YouTube(youtube_video_link)
+        except pytube_exceptions.PytubeError as e:
+            self.logger.error(
+                "An Error occurred while fetching the YouTube video: " + str(e)
+            )
+            return None
+
+        audio_stream = yt.streams.get_audio_only()
+
+        if not os.path.exists(filename):
+            try:
+                audio_stream.download(filename=filename)
+                self.logger.info(f"Downloaded {youtube_video_link} successfully.")
+            except Exception as e:
+                self.logger.error(
+                    "An Error occurred while downloading the YouTube video: " + str(e)
+                )
+                return None
+
+    def convert_to_wav(self, mp4_file: str, wav_file: str) -> None:
+        """
+        Converts a mp4 audio file into wav format.
+
+        Parameters
+        ----------
+            mp4_file : str
+                The path to the mp4 audio file.
+            wav_file : str
+                The path where the wav file should be saved
+        """
+        try:
+            sound = AudioSegment.from_file(mp4_file, format="mp4")
+            sound.export(wav_file, format="wav")
+        except Exception as e:
+            self.logger.error(
+                f"An error occurred while converting {mp4_file} to wav: {e}"
+            )
+
+    def load_local_file(self, local_file_path: str) -> str:
+        """
+        Loads a local file and converts it to wav format.
+        Return wav file path.
+
+        Parameters
+        ----------
+            local_file_path : str
+                The path to the local file.
+        """
+        file_extension = os.path.splitext(local_file_path)[-1].lower()
+
+        if file_extension in [".mp4", ".mp3"]:
+            wav_file_path = local_file_path.replace(file_extension, ".wav")
+            try:
+                audio = AudioSegment.from_file(
+                    local_file_path, format=file_extension[1:]
+                )
+                audio.export(wav_file_path, format="wav")
+            except Exception as e:
+                self.logger.error(
+                    f"An error occurred while converting {local_file_path} to wav: {e}"
+                )
+
+            return wav_file_path
+
+        elif file_extension == ".wav":
+            return local_file_path
+
+        else:
+            self.logger.error(f"Unsupported file format: {file_extension}.")
+            raise ValueError("Unsupported file format.")

tests/test_file_manager.py

@@ -1,32 +1,32 @@
 import os
 import unittest
 
-from src.file_manager import FileManager
+from src.video_manager import VideoManager
 
 
-class FileManagerTest(unittest.TestCase):
+class VideoManagerTest(unittest.TestCase):
     def setUp(self):
-        self.file_manager = FileManager()
-        self.dir = '<Your file directory>'  # replace with your directory
+        self.file_manager = VideoManager()
+        self.dir = "<Your file directory>"
 
     def test_download_youtube_audio(self):
-        youtube_link = '<Your YouTube link>'  # replace with a YouTube link
-        filename = f'{self.dir}/audio.mp4'
+        youtube_link = "<Your 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
+        os.remove(filename)
 
     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')
+        mp4_file = "<Your mp4 filename>"
+        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
+        os.remove(wav_file)
 
     def test_load_local_file(self):
-        local_file_path = '<Your mp4 filename>'  # replace with a file in your directory
+        local_file_path = "<Your mp4 filename>"
         result_path = self.file_manager.load_local_file(local_file_path)
-        self.assertEqual(result_path, local_file_path.replace('.mp4', '.wav'))
+        self.assertEqual(result_path, local_file_path.replace(".mp4", ".wav"))
 
 
 if __name__ == "__main__":

tests/test_transcriber.py

@@ -1,4 +1,3 @@
-import os
 import unittest
 
 from src.transcriber import Transcriber
@@ -9,24 +8,10 @@ class TranscriberTest(unittest.TestCase):
         self.transcriber = Transcriber()
 
     def test_transcribe(self):
-        wav_file = '<Your wav file>'  # Replace with path to a .wav file on your machine
+        wav_file = "<Your wav file>"
         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()