diff --git a/src/__init__.py b/src/__init__.py index 3dd306f..563d12a 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1 +1 @@ -"""Functions to get video data and create transcripts as well as detailed information""" \ No newline at end of file +"""Functions to get video data and create transcripts as well as detailed information""" diff --git a/src/file_manager.py b/src/file_manager.py index 6cdd0c7..c94f353 100644 --- a/src/file_manager.py +++ b/src/file_manager.py @@ -23,7 +23,9 @@ class FileManager: try: yt = YouTube(youtube_link) except pytube_exceptions.PytubeError as e: - logging.error("An Error occurred while fetching the YouTube video: " + str(e)) + logging.error( + "An Error occurred while fetching the YouTube video: " + str(e) + ) return audio_stream = yt.streams.get_audio_only() @@ -33,7 +35,9 @@ class FileManager: 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)) + logging.error( + "An Error occurred while downloading the YouTube video: " + str(e) + ) return @staticmethod @@ -61,7 +65,9 @@ class FileManager: 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}") + logging.error( + f"An error occurred while converting {local_file_path} to wav: {e}" + ) return wav_file diff --git a/src/main.py b/src/main.py index 60023fd..2c5c6f1 100644 --- a/src/main.py +++ b/src/main.py @@ -5,14 +5,12 @@ 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' - ] + input_sources = ["https://www.youtube.com/watch?v=BVVo1BaL6JA"] - hf_token = os.environ.get('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 + os.environ["HF_TOKEN"] = hf_token transcript_generator = TranscriptGenerator() transcript_generator.generate(input_sources) diff --git a/src/transcriber.py b/src/transcriber.py index ea28fc4..476ac4f 100644 --- a/src/transcriber.py +++ b/src/transcriber.py @@ -14,7 +14,12 @@ logger = logging.getLogger(__name__) class Transcriber: """A class to transcribe audio files using whisperx.""" - def __init__(self, language_code: str = "en", device: str = "cuda", compute_type: str = "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"). @@ -24,14 +29,17 @@ class Transcriber: self.device = device if torch.cuda.is_available() else "cpu" # 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) + 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 - self.hf_token = os.environ.get('HF_TOKEN') + self.hf_token = os.environ.get("HF_TOKEN") def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict: """ @@ -48,15 +56,22 @@ class Transcriber: 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) + result = whisperx.align( + result["segments"], + self.alignment_model, + self.metadata, + audio, + self.device, + 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_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) @@ -76,13 +91,13 @@ class Transcriber: """ try: # Save full results to a json file - with open(output_json_file, 'w') as f: + 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)) + 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 diff --git a/src/transcript_generator.py b/src/transcript_generator.py index e7deb55..51cd113 100644 --- a/src/transcript_generator.py +++ b/src/transcript_generator.py @@ -28,7 +28,7 @@ class TranscriptGenerator: """ for i, source in enumerate(input_sources): try: - if source.startswith('http'): + 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) @@ -36,11 +36,13 @@ class TranscriptGenerator: 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' + 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)}") + logger.exception( + f"Failed to generate transcript for source {source}. Error: {str(e)}" + ) continue - logger.info('Transcription generation completed.') + logger.info("Transcription generation completed.")