Update: Applied Black formatting

This commit is contained in:
Orihara Izaya
2023-10-09 15:56:25 +02:00
parent f92830facc
commit dc2ee55c4a
5 changed files with 48 additions and 27 deletions
+9 -3
View File
@@ -23,7 +23,9 @@ class FileManager:
try: try:
yt = YouTube(youtube_link) yt = YouTube(youtube_link)
except pytube_exceptions.PytubeError as e: 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 return
audio_stream = yt.streams.get_audio_only() audio_stream = yt.streams.get_audio_only()
@@ -33,7 +35,9 @@ class FileManager:
audio_stream.download(filename=filename) audio_stream.download(filename=filename)
logging.info(f"Downloaded {youtube_link} successfully.") logging.info(f"Downloaded {youtube_link} successfully.")
except Exception as e: 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 return
@staticmethod @staticmethod
@@ -61,7 +65,9 @@ class FileManager:
audio = AudioSegment.from_file(local_file_path, format=ext[1:]) audio = AudioSegment.from_file(local_file_path, format=ext[1:])
audio.export(wav_file, format="wav") audio.export(wav_file, format="wav")
except Exception as e: 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 return wav_file
+3 -5
View File
@@ -5,14 +5,12 @@ from src.transcript_generator import TranscriptGenerator
def main(): def main():
# input could be a YouTube link or a local file # input could be a YouTube link or a local file
input_sources = [ input_sources = ["https://www.youtube.com/watch?v=BVVo1BaL6JA"]
'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: if not hf_token:
hf_token = input("Please enter your 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 = TranscriptGenerator()
transcript_generator.generate(input_sources) transcript_generator.generate(input_sources)
+28 -13
View File
@@ -14,7 +14,12 @@ logger = logging.getLogger(__name__)
class Transcriber: class Transcriber:
"""A class to transcribe audio files using whisperx.""" """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. Init method.
:param language_code: str: Language code for whisperx transcription (default is "en"). :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" self.device = device if torch.cuda.is_available() else "cpu"
# Initialize the transcription and alignment models # Initialize the transcription and alignment models
try: try:
self.transcription_model = whisperx.load_model("large-v2", device=self.device, compute_type=compute_type) self.transcription_model = whisperx.load_model(
self.alignment_model, self.metadata = whisperx.load_align_model(language_code=language_code, "large-v2", device=self.device, compute_type=compute_type
device=self.device) )
self.alignment_model, self.metadata = whisperx.load_align_model(
language_code=language_code, device=self.device
)
except Exception as e: except Exception as e:
logger.exception("Failed to initialize models: %s", e) logger.exception("Failed to initialize models: %s", e)
raise raise
self.transcript_output = None 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: 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) result = self.transcription_model.transcribe(audio, batch_size=batch_size)
# Use the alignment model # Use the alignment model
result = whisperx.align(result["segments"], self.alignment_model, self.metadata, audio, self.device, result = whisperx.align(
return_char_alignments=False) result["segments"],
self.alignment_model,
self.metadata,
audio,
self.device,
return_char_alignments=False,
)
# Clear GPU memory # Clear GPU memory
self._clear_memory() self._clear_memory()
# Run Diarization # Run Diarization
diarize_model = whisperx.DiarizationPipeline(use_auth_token=self.hf_token, diarize_model = whisperx.DiarizationPipeline(
device=self.device) use_auth_token=self.hf_token, device=self.device
)
diarize_segments = diarize_model(audio) diarize_segments = diarize_model(audio)
result = whisperx.assign_word_speakers(diarize_segments, result) result = whisperx.assign_word_speakers(diarize_segments, result)
@@ -76,13 +91,13 @@ class Transcriber:
""" """
try: try:
# Save full results to a json file # 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) json.dump(self.transcript_output, f)
# Save just the transcripts to a .txt file # Save just the transcripts to a .txt file
transcripts = [seg['text'] for seg in self.transcript_output['segments']] transcripts = [seg["text"] for seg in self.transcript_output["segments"]]
with open(output_txt_file, 'w') as f: with open(output_txt_file, "w") as f:
f.write('\n'.join(transcripts)) f.write("\n".join(transcripts))
except Exception as e: except Exception as e:
logger.exception("Failed to save transcript: %s", e) logger.exception("Failed to save transcript: %s", e)
raise raise
+7 -5
View File
@@ -28,7 +28,7 @@ class TranscriptGenerator:
""" """
for i, source in enumerate(input_sources): for i, source in enumerate(input_sources):
try: try:
if source.startswith('http'): if source.startswith("http"):
mp4_filename = f"audio_{i + 1}.mp4" mp4_filename = f"audio_{i + 1}.mp4"
wav_filename = f"audio_{i + 1}.wav" wav_filename = f"audio_{i + 1}.wav"
self.file_manager.download_youtube_audio(source, mp4_filename) self.file_manager.download_youtube_audio(source, mp4_filename)
@@ -36,11 +36,13 @@ class TranscriptGenerator:
else: else:
wav_filename = self.file_manager.load_local_file(source) wav_filename = self.file_manager.load_local_file(source)
self.transcriber.transcribe(wav_filename) self.transcriber.transcribe(wav_filename)
json_output_file = f'result_{i + 1}.json' json_output_file = f"result_{i + 1}.json"
txt_output_file = f'transcripts_{i + 1}.txt' txt_output_file = f"transcripts_{i + 1}.txt"
self.transcriber.save_transcript(json_output_file, txt_output_file) self.transcriber.save_transcript(json_output_file, txt_output_file)
except Exception as e: 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 continue
logger.info('Transcription generation completed.') logger.info("Transcription generation completed.")