Checkout: 2023-10-10

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

README.md

@@ -22,7 +22,7 @@ sudo apt install ffmpeg
 
 To generate transcripts create a new job with your specifications (see jobs section 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:
 
@@ -36,6 +36,8 @@ On Windows (Command prompt):
 set HF_TOKEN="your-token-here"
 ```
 
+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
@@ -50,17 +52,17 @@ Jobs specify tasks for transcript creation. A job contains the following informa
 
 - `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)
 
 ## to-do
 
-- 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
+### Feature request: specific model
+
+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)
+
+### Other
+- remember to accept the conditions on models if you use huggingface
 - check warnings
 - downloader for tib-av
 - more optional setting
@@ -71,19 +73,20 @@ Jobs specify tasks for transcript creation. A job contains the following informa
 ```
 project_root/
 ├── src/
+│   ├── 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
-│   └── 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
-│   ├── models/   # models used for transcripts (are ignored by git and not necessarily)
-│   └── transcripts/  # JSON and text transcript files
+│   ├── results/   # results of the transcripts (downloaded video files, audio, JSON and text transcript files)
+│   ├── jobs/   # place to store jobs to run
+│   └── logs/  # logs of previous jobs
+├── main.py # Initial entry point for the project
 └── README.md
 ```
 
@@ -96,5 +99,4 @@ project_root/
 
 - **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/**: This directory contains all the data needed for your code as well as the results.

data/jobs/example_job.json

@@ -0,0 +1,6 @@
+{
+  "name": "myjob1",
+  "source": [
+    "https://www.youtube.com/watch?v=BVVo1BaL6JA"
+  ]
+}

main.py

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

src/file_manager.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import json
 
 from pydub import AudioSegment
 from pytube import YouTube, exceptions as pytube_exceptions
@@ -12,8 +13,8 @@ class FileManager:
     FileManager class that handles operations such as downloading and converting audio.
     """
 
-    def __init__(self):
-        pass
+    def __init__(self, folder: str="data/results"):
+        self.folder = folder
 
     @staticmethod
     def download_youtube_audio(youtube_link: str, filename: str):
@@ -77,3 +78,26 @@ class FileManager:
         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/transcriber.py

@@ -1,5 +1,4 @@
 import gc
-import json
 import logging
 import os
 
@@ -83,25 +82,6 @@ class Transcriber:
             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():
         """

src/transcript_generator.py

@@ -15,9 +15,10 @@ class TranscriptGenerator:
     to generate transcriptions from given input sources.
     """
 
-    def __init__(self):
+    def __init__(self, job_name, language_code="en"):
         self.file_manager = FileManager()
-        self.transcriber = Transcriber()
+        self.transcriber = Transcriber(language_code)
+        self.job_name = job_name
 
     def generate(self, input_sources: List[str]) -> None:
         """
@@ -29,16 +30,16 @@ class TranscriptGenerator:
         for i, source in enumerate(input_sources):
             try:
                 if source.startswith("http"):
-                    mp4_filename = f"audio_{i + 1}.mp4"
-                    wav_filename = f"audio_{i + 1}.wav"
+                    mp4_filename = f"audio_{self.job_name}_{i + 1}.mp4"
+                    wav_filename = f"audio_{self.job_name}_{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)
+                transcript_data = self.transcriber.transcribe(wav_filename)
+                json_output_file = f"result_{self.job_name}_{i + 1}.json"
+                txt_output_file = f"transcripts_{self.job_name}_{i + 1}.txt"
+                self.file_manager.save_transcript(transcript_data, json_output_file, txt_output_file)
             except Exception as e:
                 logger.exception(
                     f"Failed to generate transcript for source {source}. Error: {str(e)}"