Update: added script to handle data and transcribe.

This commit is contained in:
2023-09-20 15:07:25 +02:00
parent dca32539cf
commit 12065e3124
5 changed files with 73 additions and 0 deletions
+14
View File
@@ -0,0 +1,14 @@
from src.data_handling import download_audio, convert_to_wav
from src.transcriber import transcribe_wav_file
youtube_links = ['https://www.youtube.com/watch?v=BVVo1BaL6JA', 'https://www.youtube.com/watch?v=m0gr3snccVY&list=PL8p82vzFeTAj3cbHaXE9sC48pm4oZY2rH&index=2', 'https://www.youtube.com/watch?v=ThjgUfck8Sc&list=PL8p82vzFeTAj3cbHaXE9sC48pm4oZY2rH&index=3']
#youtube_links = ['https://www.youtube.com/watch?v=6A5ZSiosgAk']
for i, link in enumerate(youtube_links):
mp4_filename = f"audio_{i+1}.mp4"
wav_filename = f"audio_{i+1}.wav"
download_audio(link, mp4_filename)
convert_to_wav(mp4_filename, wav_filename)
transcribe_wav_file(wav_filename, i)
+5
View File
@@ -0,0 +1,5 @@
youtube-dl
pfy
pytube
pydub
git+https://github.com/m-bain/whisperx.git
+1
View File
@@ -0,0 +1 @@
"""Functions to get video data and create transcripts as well as detailed information"""
+13
View File
@@ -0,0 +1,13 @@
def download_audio(youtube_link, filename):
yt = YouTube(youtube_link)
# get highest quality audio stream
audio_stream = yt.streams.get_audio_only()
# download
if not os.path.exists(filename):
audio_stream.download(filename=filename)
def convert_to_wav(mp4_file, wav_file):
sound = AudioSegment.from_file(mp4_file, format="mp4")
sound.export(wav_file, format="wav")
+40
View File
@@ -0,0 +1,40 @@
google_colab_gpu = "cuda" if torch.cuda.is_available() else "cpu"
def transcribe_wav_file(wav_filename, index):
batch_size = 16
compute_type = "float16"
model = whisperx.load_model("large-v2", device=google_colab_gpu, compute_type=compute_type)
audio = whisperx.load_audio(wav_filename)
result = model.transcribe(audio, batch_size=batch_size)
model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=google_colab_gpu)
result = whisperx.align(result["segments"], model_a, metadata, audio, google_colab_gpu, return_char_alignments=False)
# clear GPU memory
del model
del model_a
del metadata
gc.collect()
torch.cuda.empty_cache()
diarize_model = whisperx.DiarizationPipeline(use_auth_token='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', device=google_colab_gpu)
diarize_segments = diarize_model(audio)
result = whisperx.assign_word_speakers(diarize_segments, result)
# clear GPU memory
del diarize_model
del diarize_segments
gc.collect()
torch.cuda.empty_cache()
# Save full results to a json file
with open(f'result_{index+1}.json', 'w') as f:
json.dump(result, f)
# Save just the transcripts to a .txt file
transcripts = [seg['value'] for seg in result['segments']]
with open(f'transcripts_{index+1}.txt', 'w') as f:
f.write('\n'.join(transcripts))