Checkout: 2023-10-10

save in data folder and refactor save_transcript to be part of FileManager class

.gitignore

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

LICENSE

@@ -1,21 +0,0 @@
-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,27 +1,15 @@
-# Project_whisperTrans_Code
+# Project_whisperTrans01_Code
 
-## Overview
+## Content
 
-whisperTrans is a video2transcript tool built upon [Whisper](https://github.com/openai/whisper)
+Transcriber code
-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
 
-The project is tested under python 3.10.12.
+Before running this script, make sure to install the necessary python packages using the following command:
-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  --index-url https://download.pytorch.org/whl/cu118
+pip install -r requirements.txt
 ```
 
 If your system does not have ffmpeg installed you can install it via:
@@ -30,46 +18,11 @@ If your system does not have ffmpeg installed you can install it via:
 sudo apt install ffmpeg
 ```
 
-For GPU usage you need to install the CUDA toolkit. You can install it via:
+## Running the script
 
-```bash
+To generate transcripts create a new job with your specifications (see jobs section below).
-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:
+If you need to download the transcriber model set your huggingface token (initially running or using new language):
 
 On Linux/Mac:
 
@@ -83,72 +36,67 @@ On Windows (Command prompt):
 set HF_TOKEN="your-token-here"
 ```
 
-After setting up the token, execute the script with:
+If the token is not set you will be asked for it when running the software
+
+Then run the following command to generate your transcript:
 
 ```bash
-python main.py --job {name}
+main.py --job {name}
 ```
 
-Here, `name` is the identifier for your transcription job.
+where name is the name of your 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)
+- `source`: List[string]: The sources of the videos (paths, YouTube links, TibAV links)
-- `language`: Optional[string]: The language of the video (If not specified employs auto-detection) as a language code
+- `language`: Optional[string]: The language of the video (If not specified employs auto-detection)
-  (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
 
-## Contributing
+## to-do
 
-Contributions are welcome! Please fork the repository and make your changes in a separate branch. Submit a pull request
+### Feature request: specific model
-with your proposed modifications. For major changes, please open an issue first to discuss what you would like to
-change.
 
-## License
+add an opitonal model path to the json config
+- `model_path`:Optional[string]: The path to a specific language model to be used (mainly for testing)
 
-This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
+### Other
-
+- remember to accept the conditions on models if you use huggingface
-## Contact
+- check warnings
-
+- downloader for tib-av
-This project is maintained by [Gerrit Anders](https://www.iwm-tuebingen.de/www/de/personen/ma.html?uid=ganders).
+- more optional setting
-For any questions or support, please open an issue on this repository or reach out via email
+- add progress bar
-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/
-│   ├── video_manager.py   # Contains VideoManager class
+│   ├── init.py # init so the main can load all necessary classes
+│   ├── 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
+│   ├── data_handling.py  # Contains download_audio and convert_to_wav methods
 ├── tests/    
 │   ├── test_file_manager.py
 │   ├── test_transcriber.py
 │   └── test_transcript_generator.py
 ├── data/
-│   ├── logs/  # logs of previous jobs
+│   ├── results/   # results of the transcripts (downloaded video files, audio, JSON and text transcript files)
-│   ├── jobs/   # place to store job files
+│   ├── jobs/   # place to store jobs to run
-│   └── results/   # place to store results
+│   └── logs/  # logs of previous jobs
-├── main.py    # Main script that runs the process
+├── main.py # Initial entry point for the project
-├── 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 as well as the results.

data/jobs/example_job.json

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

main.py

@@ -1,50 +1,36 @@
-import json
-import logging
 import os
-from argparse import ArgumentParser, Namespace
+import json
-from pathlib import Path
+import argparse
 
-from src.logger import setup_logger
 from src.transcript_generator import TranscriptGenerator
 
 
 def main():
-    parser: ArgumentParser = ArgumentParser(
+    # Create a parser
-        description="Transcript Generator for Videos"
+    parser = argparse.ArgumentParser(description='Transcript Generator for Videos')
-    )
 
-    parser.add_argument("--job", type=str, help="The name of the job", required=True)
+    # Add the arguments
+    parser.add_argument('--job', type=str, help='The name of the job', required=True)
 
-    arguments: Namespace = parser.parse_args()
+    # Parse the arguments
+    args = parser.parse_args()
 
-    with open(f"data/jobs/{arguments.job}", "r") as f:
+    # Load the job config
+    with open(f'jobs/{args.job}.json', 'r') as f:
         job_config = json.load(f)
 
-    log_directory: Path = Path("data/logs")
+    # Check the hf_token
-    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 = TranscriptGenerator(
+    transcript_generator = TranscriptGenerator(
-        job_name=job_config["name"],
+        name=job_config['name'],
-        language_code=job_config.get("language_code", "en"),
+        language_code=job_config['language_code']
-        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"])
+    transcript_generator.generate(job_config['source'])
 
-
+    if __name__ == "__main__":
-if __name__ == "__main__":
+        main()
-    main()

requirements.txt

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

src/file_manager.py

@@ -0,0 +1,103 @@
+import logging
+import os
+import json
+
+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, folder: str="data/results"):
+        self.folder = folder
+
+    @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.")
+
+    def save_transcript(self, transcript_data: dict, output_json_file: str, output_txt_file: str):
+        """
+        Save the transcription result to a json file and a text file.
+        :param transcript_data: dict: Transcription data.
+        :param output_json_file: str: Name for the output json file.
+        :param output_txt_file: str: Name for the output text file.
+        """
+        output_json_file = os.path.join(self.folder, output_json_file)
+        output_txt_file = os.path.join(self.folder, output_txt_file)
+
+        try:
+            # Save full results to a json file
+            with open(output_json_file, "w") as f:
+                json.dump(transcript_data, f)
+
+            # Save just the transcripts to a .txt file
+            transcripts = [seg["text"] for seg in transcript_data["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

src/logger.py

@@ -1,47 +0,0 @@
-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

@@ -1,128 +0,0 @@
-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

@@ -5,57 +5,37 @@ 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,
     ):
-        if asr_options is None:
+        """
-            asr_options = {
+        Init method.
-                "repetition_penalty": 1,
+        :param language_code: str: Language code for whisperx transcription (default is "en").
-                "prompt_reset_on_temperature": 0.5,
+        :param device: str: Device to use for transcription (default is "cuda" if available, else "cpu").
-                "no_repeat_ngram_size": 0,
+        :param compute_type: str: Compute type for whisperx transcription (default is "float16").
-            }
+        """
-        self.logger = logging.getLogger(__name__)
-
         self.device = device if torch.cuda.is_available() else "cpu"
-        self.logger.info("Using device: %s", self.device)
+        # Initialize the transcription and alignment models
-
         try:
             self.transcription_model = whisperx.load_model(
-                "large-v2",
+                "large-v2", device=self.device, compute_type=compute_type
-                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:
-            self.logger.exception("Failed to initialize models: %s", e)
+            logger.exception("Failed to initialize models: %s", e)
             raise
         self.transcript_output = None
         self.hf_token = os.environ.get("HF_TOKEN")
@@ -63,24 +43,18 @@ 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.
-        Parameters:
+        :param batch_size: int: Batch size for whisperx transcription (default is 16).
-        ----------
+        :return: dict: Transcribed and aligned result.
-        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,
@@ -90,22 +64,26 @@ 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:
-            self.logger.exception("Failed to transcribe: %s", e)
+            logger.exception("Failed to transcribe: %s", e)
             raise
 
-    def _clear_memory(self):
+    @staticmethod
+    def _clear_memory():
         """
         Clear GPU memory.
         """
@@ -113,5 +91,5 @@ class Transcriber:
             gc.collect()
             torch.cuda.empty_cache()
         except Exception as e:
-            self.logger.exception("Failed to clear GPU memory: %s", e)
+            logger.exception("Failed to clear GPU memory: %s", e)
             raise

src/transcript_generator.py

@@ -1,82 +1,49 @@
 import logging
+from typing import List
 
-from src.output_manager import OutputManager
+from src.file_manager import FileManager
 from src.transcriber import Transcriber
-from src.video_manager import VideoManager
+
+# 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 in desired formats from given input sources.
+    to generate transcriptions 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__(
+    def __init__(self, job_name, language_code="en"):
-        self,
+        self.file_manager = FileManager()
-        job_name: str,
+        self.transcriber = Transcriber(language_code)
-        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
 
-        self.logger = logging.getLogger(__name__)
+    def generate(self, input_sources: List[str]) -> None:
-
-    def generate(self, input_sources: list[str]) -> None:
         """
         Generates transcriptions for the given list of input sources.
 
-
+        Args:
-        Parameters
+            input_sources (List[str]): List of input sources which could be URLs or local file paths.
-        ----------
-            input_sources : list[str]
-                A list of input sources (local file paths or YouTube URLs) for which transcriptions should be generated.
         """
-        for source_number, source in enumerate(input_sources):
+        for i, source in enumerate(input_sources):
             try:
                 if source.startswith("http"):
-                    mp4_filename = f"audio_{self.job_name}_{source_number + 1}.mp4"
+                    mp4_filename = f"audio_{self.job_name}_{i + 1}.mp4"
-                    wav_filename = f"audio_{self.job_name}_{source_number + 1}.wav"
+                    wav_filename = f"audio_{self.job_name}_{i + 1}.wav"
-                    self.video_manager.download_youtube_audio(source, mp4_filename)
+                    self.file_manager.download_youtube_audio(source, mp4_filename)
-                    self.video_manager.convert_to_wav(mp4_filename, wav_filename)
+                    self.file_manager.convert_to_wav(mp4_filename, wav_filename)
                 else:
-                    wav_filename = self.video_manager.load_local_file(source)
+                    wav_filename = self.file_manager.load_local_file(source)
                 transcript_data = self.transcriber.transcribe(wav_filename)
-                self.output_manager.save_transcript(
+                json_output_file = f"result_{self.job_name}_{i + 1}.json"
-                    transcript_data, self.job_name, source_number
+                txt_output_file = f"transcripts_{self.job_name}_{i + 1}.txt"
-                )
+                self.file_manager.save_transcript(transcript_data, json_output_file, txt_output_file)
-
-                self.logger.info(f"Transcription generated for source {source}.")
-
             except Exception as e:
-                self.logger.exception(
+                logger.exception(
                     f"Failed to generate transcript for source {source}. Error: {str(e)}"
                 )
                 continue
 
-        self.logger.info("Transcription generation completed.")
+        logger.info("Transcription generation completed.")

src/video_manager.py

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

tests/test_transcriber.py

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