From 57cc73a4395cd090d5d3ebf23990f243b0d71617 Mon Sep 17 00:00:00 2001 From: Izaya Orihara Date: Fri, 6 Oct 2023 15:45:14 +0200 Subject: [PATCH] Update: Refactored the code --- README.md | 5 -- src/file_manager.py | 70 +++++++++++++++------ src/transcriber.py | 121 +++++++++++++++++++++++------------- src/transcript_generator.py | 51 ++++++++++----- 4 files changed, 165 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index 8bf4ba0..0fb7ac8 100644 --- a/README.md +++ b/README.md @@ -4,13 +4,8 @@ Transcriber code -Currently very unorganized - ## to-do -- clean the code -- organize the stuff -- add functions for other data sources - seperate data production form code ``` diff --git a/src/file_manager.py b/src/file_manager.py index 40755eb..6cdd0c7 100644 --- a/src/file_manager.py +++ b/src/file_manager.py @@ -1,41 +1,73 @@ +import logging import os from pydub import AudioSegment -from pytube import YouTube +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 - def download_youtube_audio(self, youtube_link, filename): - yt = YouTube(youtube_link) + @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 - # Get the highest quality audio stream audio_stream = yt.streams.get_audio_only() - # Download if the file doesn't already exist if not os.path.exists(filename): - audio_stream.download(filename=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 - def convert_to_wav(self, mp4_file, wav_file): - sound = AudioSegment.from_file(mp4_file, format="mp4") - sound.export(wav_file, format="wav") + @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}") - def load_local_file(self, local_file_path): - # Check the file type + @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 == ".mp4": - wav_file = local_file_path.replace(".mp4", ".wav") - self.convert_to_wav(local_file_path, wav_file) - return wav_file - elif ext == ".mp3": - wav_file = local_file_path.replace(".mp3", ".wav") - audio = AudioSegment.from_file(local_file_path, format="mp3") - audio.export(wav_file, format="wav") + 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.") diff --git a/src/transcriber.py b/src/transcriber.py index b13232f..d73592b 100644 --- a/src/transcriber.py +++ b/src/transcriber.py @@ -1,63 +1,100 @@ import gc import json +import logging 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.""" - def __init__(self, language_code="en", device="cuda", compute_type="float16"): + def __init__(self, language_code: str = "en", device: str = "cuda", compute_type: str = "float16"): + """ + 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"). + """ self.device = device if torch.cuda.is_available() else "cpu" - - # Initialize the transcription model - self.transcription_model = whisperx.load_model( - "large-v2", device=self.device, compute_type=compute_type) - - # Initialize the alignment model - self.alignment_model, self.metadata = whisperx.load_align_model( - language_code=language_code, device=self.device) + # Initialize the transcription and alignment models + try: + self.transcription_model = whisperx.load_model("large-v2", device=self.device, compute_type=compute_type) + 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) + raise self.transcript_output = None - def transcribe(self, wav_filename, batch_size=16): - # Load the audio file - audio = whisperx.load_audio(wav_filename) + 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. + """ + 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) + # 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, self.metadata, audio, self.device, - return_char_alignments=False) + # Use the alignment model + result = whisperx.align(result["segments"], self.alignment_model, self.metadata, audio, self.device, + return_char_alignments=False) - # Clear GPU memory - self._clear_memory() + # Clear GPU memory + self._clear_memory() - # Run Diarization - diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', - device=self.device) - diarize_segments = diarize_model(audio) - result = whisperx.assign_word_speakers(diarize_segments, result) + # Run Diarization + diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', + 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 + # Clear GPU memory again + self._clear_memory() + self.transcript_output = result + return result + except Exception as e: + logger.exception("Failed to transcribe: %s", e) + raise - def save_transcript(self, output_json_file, output_txt_file): - # Save full results to a json file - with open(output_json_file, 'w') as f: - json.dump(self.transcript_output, f) + 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['value'] for seg in self.transcript_output['segments']] - with open(output_txt_file, 'w') as f: - f.write('\n'.join(transcripts)) + # Save just the transcripts to a .txt file + transcripts = [seg['value'] 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 def _clear_memory(self): - del self.transcription_model - del self.alignment_model - del self.metadata - gc.collect() - torch.cuda.empty_cache() + """ + Clear GPU memory. + """ + try: + del self.transcription_model + del self.alignment_model + del self.metadata + gc.collect() + torch.cuda.empty_cache() + except Exception as e: + logger.exception("Failed to clear GPU memory: %s", e) + raise diff --git a/src/transcript_generator.py b/src/transcript_generator.py index 79343cb..e7deb55 100644 --- a/src/transcript_generator.py +++ b/src/transcript_generator.py @@ -1,27 +1,46 @@ +import logging +from typing import List + from src.file_manager import FileManager from src.transcriber import Transcriber +# 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 from given input sources. + """ + def __init__(self): self.file_manager = FileManager() self.transcriber = Transcriber() - def generate(self, input_sources): + 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. + """ for i, source in enumerate(input_sources): - if source.startswith('http'): - mp4_filename = f"audio_{i + 1}.mp4" - wav_filename = f"audio_{i + 1}.wav" + 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) + 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) + except Exception as e: + logger.exception(f"Failed to generate transcript for source {source}. Error: {str(e)}") + continue - 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) - - transcription = 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) - - print('Transcription generation completed.') + logger.info('Transcription generation completed.')