diff --git a/README.md b/README.md index d7ebe41..9e5c48f 100644 --- a/README.md +++ b/README.md @@ -56,13 +56,17 @@ Jobs specify tasks for transcript creation. A job contains the following informa ## to-do -### Feature: specific model +### 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) +### Refactoring: saving transcripts + +let the generate transcript return a transcript object that is the saved by file manager +give file manager an init to specify the saving folder + ### Other -- seperate data production form code - remember to accept the conditions on models if you use huggingface - check warnings - downloader for tib-av diff --git a/src/file_manager.py b/src/file_manager.py index c94f353..ba5c0e8 100644 --- a/src/file_manager.py +++ b/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 diff --git a/src/transcriber.py b/src/transcriber.py index 4842b7a..91a4378 100644 --- a/src/transcriber.py +++ b/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(): """ diff --git a/src/transcript_generator.py b/src/transcript_generator.py index 03594d3..3229cec 100644 --- a/src/transcript_generator.py +++ b/src/transcript_generator.py @@ -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)}"