Refactoring: saving transcripts

save in data folder and refactor save_transcript to be part of FileManager class
This commit is contained in:
2023-10-10 15:36:16 +02:00
parent e2b5ad15b4
commit 29df782f05
4 changed files with 34 additions and 26 deletions
+26 -2
View File
@@ -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
-20
View File
@@ -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():
"""
+2 -2
View File
@@ -36,10 +36,10 @@ class TranscriptGenerator:
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)
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.transcriber.save_transcript(json_output_file, txt_output_file)
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)}"