Kye
1 year ago
parent
7cd6f25353
commit
d33588becc
@ -1,138 +0,0 @@
|
|||||||
import os
|
|
||||||
import subprocess
|
|
||||||
|
|
||||||
try:
|
|
||||||
import whisperx
|
|
||||||
from pydub import AudioSegment
|
|
||||||
from pytube import YouTube
|
|
||||||
except Exception as error:
|
|
||||||
print("Error importing pytube. Please install pytube manually.")
|
|
||||||
print("pip install pytube")
|
|
||||||
print("pip install pydub")
|
|
||||||
print("pip install whisperx")
|
|
||||||
print(f"Pytube error: {error}")
|
|
||||||
|
|
||||||
|
|
||||||
class WhisperX:
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
video_url,
|
|
||||||
audio_format="mp3",
|
|
||||||
device="cuda",
|
|
||||||
batch_size=16,
|
|
||||||
compute_type="float16",
|
|
||||||
hf_api_key=None,
|
|
||||||
):
|
|
||||||
"""
|
|
||||||
# Example usage
|
|
||||||
video_url = "url"
|
|
||||||
speech_to_text = WhisperX(video_url)
|
|
||||||
transcription = speech_to_text.transcribe_youtube_video()
|
|
||||||
print(transcription)
|
|
||||||
|
|
||||||
"""
|
|
||||||
self.video_url = video_url
|
|
||||||
self.audio_format = audio_format
|
|
||||||
self.device = device
|
|
||||||
self.batch_size = batch_size
|
|
||||||
self.compute_type = compute_type
|
|
||||||
self.hf_api_key = hf_api_key
|
|
||||||
|
|
||||||
def install(self):
|
|
||||||
subprocess.run(["pip", "install", "whisperx"])
|
|
||||||
subprocess.run(["pip", "install", "pytube"])
|
|
||||||
subprocess.run(["pip", "install", "pydub"])
|
|
||||||
|
|
||||||
def download_youtube_video(self):
|
|
||||||
audio_file = f"video.{self.audio_format}"
|
|
||||||
|
|
||||||
# Download video 📥
|
|
||||||
yt = YouTube(self.video_url)
|
|
||||||
yt_stream = yt.streams.filter(only_audio=True).first()
|
|
||||||
yt_stream.download(filename="video.mp4")
|
|
||||||
|
|
||||||
# Convert video to audio 🎧
|
|
||||||
video = AudioSegment.from_file("video.mp4", format="mp4")
|
|
||||||
video.export(audio_file, format=self.audio_format)
|
|
||||||
os.remove("video.mp4")
|
|
||||||
|
|
||||||
return audio_file
|
|
||||||
|
|
||||||
def transcribe_youtube_video(self):
|
|
||||||
audio_file = self.download_youtube_video()
|
|
||||||
|
|
||||||
device = "cuda"
|
|
||||||
batch_size = 16
|
|
||||||
compute_type = "float16"
|
|
||||||
|
|
||||||
# 1. Transcribe with original Whisper (batched) 🗣️
|
|
||||||
model = whisperx.load_model(
|
|
||||||
"large-v2", device, compute_type=compute_type
|
|
||||||
)
|
|
||||||
audio = whisperx.load_audio(audio_file)
|
|
||||||
result = model.transcribe(audio, batch_size=batch_size)
|
|
||||||
|
|
||||||
# 2. Align Whisper output 🔍
|
|
||||||
model_a, metadata = whisperx.load_align_model(
|
|
||||||
language_code=result["language"], device=device
|
|
||||||
)
|
|
||||||
result = whisperx.align(
|
|
||||||
result["segments"],
|
|
||||||
model_a,
|
|
||||||
metadata,
|
|
||||||
audio,
|
|
||||||
device,
|
|
||||||
return_char_alignments=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
# 3. Assign speaker labels 🏷️
|
|
||||||
diarize_model = whisperx.DiarizationPipeline(
|
|
||||||
use_auth_token=self.hf_api_key, device=device
|
|
||||||
)
|
|
||||||
diarize_model(audio_file)
|
|
||||||
|
|
||||||
try:
|
|
||||||
segments = result["segments"]
|
|
||||||
transcription = " ".join(
|
|
||||||
segment["text"] for segment in segments
|
|
||||||
)
|
|
||||||
return transcription
|
|
||||||
except KeyError:
|
|
||||||
print("The key 'segments' is not found in the result.")
|
|
||||||
|
|
||||||
def transcribe(self, audio_file):
|
|
||||||
model = whisperx.load_model(
|
|
||||||
"large-v2", self.device, self.compute_type
|
|
||||||
)
|
|
||||||
audio = whisperx.load_audio(audio_file)
|
|
||||||
result = model.transcribe(audio, batch_size=self.batch_size)
|
|
||||||
|
|
||||||
# 2. Align Whisper output 🔍
|
|
||||||
model_a, metadata = whisperx.load_align_model(
|
|
||||||
language_code=result["language"], device=self.device
|
|
||||||
)
|
|
||||||
|
|
||||||
result = whisperx.align(
|
|
||||||
result["segments"],
|
|
||||||
model_a,
|
|
||||||
metadata,
|
|
||||||
audio,
|
|
||||||
self.device,
|
|
||||||
return_char_alignments=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
# 3. Assign speaker labels 🏷️
|
|
||||||
diarize_model = whisperx.DiarizationPipeline(
|
|
||||||
use_auth_token=self.hf_api_key, device=self.device
|
|
||||||
)
|
|
||||||
|
|
||||||
diarize_model(audio_file)
|
|
||||||
|
|
||||||
try:
|
|
||||||
segments = result["segments"]
|
|
||||||
transcription = " ".join(
|
|
||||||
segment["text"] for segment in segments
|
|
||||||
)
|
|
||||||
return transcription
|
|
||||||
except KeyError:
|
|
||||||
print("The key 'segments' is not found in the result.")
|
|
||||||
@ -1,222 +0,0 @@
|
|||||||
import os
|
|
||||||
import subprocess
|
|
||||||
import tempfile
|
|
||||||
from unittest.mock import patch
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
import whisperx
|
|
||||||
from pydub import AudioSegment
|
|
||||||
from pytube import YouTube
|
|
||||||
from swarms.models.whisperx_model import WhisperX
|
|
||||||
|
|
||||||
|
|
||||||
# Fixture to create a temporary directory for testing
|
|
||||||
@pytest.fixture
|
|
||||||
def temp_dir():
|
|
||||||
with tempfile.TemporaryDirectory() as tempdir:
|
|
||||||
yield tempdir
|
|
||||||
|
|
||||||
|
|
||||||
# Mock subprocess.run to prevent actual installation during tests
|
|
||||||
@patch.object(subprocess, "run")
|
|
||||||
def test_speech_to_text_install(mock_run):
|
|
||||||
stt = WhisperX("https://www.youtube.com/watch?v=MJd6pr16LRM")
|
|
||||||
stt.install()
|
|
||||||
mock_run.assert_called_with(["pip", "install", "whisperx"])
|
|
||||||
|
|
||||||
|
|
||||||
# Mock pytube.YouTube and pytube.Streams for download tests
|
|
||||||
@patch("pytube.YouTube")
|
|
||||||
@patch.object(YouTube, "streams")
|
|
||||||
def test_speech_to_text_download_youtube_video(
|
|
||||||
mock_streams, mock_youtube, temp_dir
|
|
||||||
):
|
|
||||||
# Mock YouTube and streams
|
|
||||||
video_url = "https://www.youtube.com/watch?v=MJd6pr16LRM"
|
|
||||||
mock_stream = mock_streams().filter().first()
|
|
||||||
mock_stream.download.return_value = os.path.join(
|
|
||||||
temp_dir, "video.mp4"
|
|
||||||
)
|
|
||||||
mock_youtube.return_value = mock_youtube
|
|
||||||
mock_youtube.streams = mock_streams
|
|
||||||
|
|
||||||
stt = WhisperX(video_url)
|
|
||||||
audio_file = stt.download_youtube_video()
|
|
||||||
|
|
||||||
assert os.path.exists(audio_file)
|
|
||||||
assert audio_file.endswith(".mp3")
|
|
||||||
|
|
||||||
|
|
||||||
# Mock whisperx.load_model and whisperx.load_audio for transcribe tests
|
|
||||||
@patch("whisperx.load_model")
|
|
||||||
@patch("whisperx.load_audio")
|
|
||||||
@patch("whisperx.load_align_model")
|
|
||||||
@patch("whisperx.align")
|
|
||||||
@patch.object(whisperx.DiarizationPipeline, "__call__")
|
|
||||||
def test_speech_to_text_transcribe_youtube_video(
|
|
||||||
mock_diarization,
|
|
||||||
mock_align,
|
|
||||||
mock_align_model,
|
|
||||||
mock_load_audio,
|
|
||||||
mock_load_model,
|
|
||||||
temp_dir,
|
|
||||||
):
|
|
||||||
# Mock whisperx functions
|
|
||||||
mock_load_model.return_value = mock_load_model
|
|
||||||
mock_load_model.transcribe.return_value = {
|
|
||||||
"language": "en",
|
|
||||||
"segments": [{"text": "Hello, World!"}],
|
|
||||||
}
|
|
||||||
|
|
||||||
mock_load_audio.return_value = "audio_path"
|
|
||||||
mock_align_model.return_value = (mock_align_model, "metadata")
|
|
||||||
mock_align.return_value = {
|
|
||||||
"segments": [{"text": "Hello, World!"}]
|
|
||||||
}
|
|
||||||
|
|
||||||
# Mock diarization pipeline
|
|
||||||
mock_diarization.return_value = None
|
|
||||||
|
|
||||||
video_url = "https://www.youtube.com/watch?v=MJd6pr16LRM/video"
|
|
||||||
stt = WhisperX(video_url)
|
|
||||||
transcription = stt.transcribe_youtube_video()
|
|
||||||
|
|
||||||
assert transcription == "Hello, World!"
|
|
||||||
|
|
||||||
|
|
||||||
# More tests for different scenarios and edge cases can be added here.
|
|
||||||
|
|
||||||
|
|
||||||
# Test transcribe method with provided audio file
|
|
||||||
def test_speech_to_text_transcribe_audio_file(temp_dir):
|
|
||||||
# Create a temporary audio file
|
|
||||||
audio_file = os.path.join(temp_dir, "test_audio.mp3")
|
|
||||||
AudioSegment.silent(duration=500).export(audio_file, format="mp3")
|
|
||||||
|
|
||||||
stt = WhisperX("https://www.youtube.com/watch?v=MJd6pr16LRM")
|
|
||||||
transcription = stt.transcribe(audio_file)
|
|
||||||
|
|
||||||
assert transcription == ""
|
|
||||||
|
|
||||||
|
|
||||||
# Test transcribe method when Whisperx fails
|
|
||||||
@patch("whisperx.load_model")
|
|
||||||
@patch("whisperx.load_audio")
|
|
||||||
def test_speech_to_text_transcribe_whisperx_failure(
|
|
||||||
mock_load_audio, mock_load_model, temp_dir
|
|
||||||
):
|
|
||||||
# Mock whisperx functions to raise an exception
|
|
||||||
mock_load_model.side_effect = Exception("Whisperx failed")
|
|
||||||
mock_load_audio.return_value = "audio_path"
|
|
||||||
|
|
||||||
stt = WhisperX("https://www.youtube.com/watch?v=MJd6pr16LRM")
|
|
||||||
transcription = stt.transcribe("audio_path")
|
|
||||||
|
|
||||||
assert transcription == "Whisperx failed"
|
|
||||||
|
|
||||||
|
|
||||||
# Test transcribe method with missing 'segments' key in Whisperx output
|
|
||||||
@patch("whisperx.load_model")
|
|
||||||
@patch("whisperx.load_audio")
|
|
||||||
@patch("whisperx.load_align_model")
|
|
||||||
@patch("whisperx.align")
|
|
||||||
@patch.object(whisperx.DiarizationPipeline, "__call__")
|
|
||||||
def test_speech_to_text_transcribe_missing_segments(
|
|
||||||
mock_diarization,
|
|
||||||
mock_align,
|
|
||||||
mock_align_model,
|
|
||||||
mock_load_audio,
|
|
||||||
mock_load_model,
|
|
||||||
):
|
|
||||||
# Mock whisperx functions to return incomplete output
|
|
||||||
mock_load_model.return_value = mock_load_model
|
|
||||||
mock_load_model.transcribe.return_value = {"language": "en"}
|
|
||||||
|
|
||||||
mock_load_audio.return_value = "audio_path"
|
|
||||||
mock_align_model.return_value = (mock_align_model, "metadata")
|
|
||||||
mock_align.return_value = {}
|
|
||||||
|
|
||||||
# Mock diarization pipeline
|
|
||||||
mock_diarization.return_value = None
|
|
||||||
|
|
||||||
stt = WhisperX("https://www.youtube.com/watch?v=MJd6pr16LRM")
|
|
||||||
transcription = stt.transcribe("audio_path")
|
|
||||||
|
|
||||||
assert transcription == ""
|
|
||||||
|
|
||||||
|
|
||||||
# Test transcribe method with Whisperx align failure
|
|
||||||
@patch("whisperx.load_model")
|
|
||||||
@patch("whisperx.load_audio")
|
|
||||||
@patch("whisperx.load_align_model")
|
|
||||||
@patch("whisperx.align")
|
|
||||||
@patch.object(whisperx.DiarizationPipeline, "__call__")
|
|
||||||
def test_speech_to_text_transcribe_align_failure(
|
|
||||||
mock_diarization,
|
|
||||||
mock_align,
|
|
||||||
mock_align_model,
|
|
||||||
mock_load_audio,
|
|
||||||
mock_load_model,
|
|
||||||
):
|
|
||||||
# Mock whisperx functions to raise an exception during align
|
|
||||||
mock_load_model.return_value = mock_load_model
|
|
||||||
mock_load_model.transcribe.return_value = {
|
|
||||||
"language": "en",
|
|
||||||
"segments": [{"text": "Hello, World!"}],
|
|
||||||
}
|
|
||||||
|
|
||||||
mock_load_audio.return_value = "audio_path"
|
|
||||||
mock_align_model.return_value = (mock_align_model, "metadata")
|
|
||||||
mock_align.side_effect = Exception("Align failed")
|
|
||||||
|
|
||||||
# Mock diarization pipeline
|
|
||||||
mock_diarization.return_value = None
|
|
||||||
|
|
||||||
stt = WhisperX("https://www.youtube.com/watch?v=MJd6pr16LRM")
|
|
||||||
transcription = stt.transcribe("audio_path")
|
|
||||||
|
|
||||||
assert transcription == "Align failed"
|
|
||||||
|
|
||||||
|
|
||||||
# Test transcribe_youtube_video when Whisperx diarization fails
|
|
||||||
@patch("pytube.YouTube")
|
|
||||||
@patch.object(YouTube, "streams")
|
|
||||||
@patch("whisperx.DiarizationPipeline")
|
|
||||||
@patch("whisperx.load_audio")
|
|
||||||
@patch("whisperx.load_align_model")
|
|
||||||
@patch("whisperx.align")
|
|
||||||
def test_speech_to_text_transcribe_diarization_failure(
|
|
||||||
mock_align,
|
|
||||||
mock_align_model,
|
|
||||||
mock_load_audio,
|
|
||||||
mock_diarization,
|
|
||||||
mock_streams,
|
|
||||||
mock_youtube,
|
|
||||||
temp_dir,
|
|
||||||
):
|
|
||||||
# Mock YouTube and streams
|
|
||||||
video_url = "https://www.youtube.com/watch?v=MJd6pr16LRM"
|
|
||||||
mock_stream = mock_streams().filter().first()
|
|
||||||
mock_stream.download.return_value = os.path.join(
|
|
||||||
temp_dir, "video.mp4"
|
|
||||||
)
|
|
||||||
mock_youtube.return_value = mock_youtube
|
|
||||||
mock_youtube.streams = mock_streams
|
|
||||||
|
|
||||||
# Mock whisperx functions
|
|
||||||
mock_load_audio.return_value = "audio_path"
|
|
||||||
mock_align_model.return_value = (mock_align_model, "metadata")
|
|
||||||
mock_align.return_value = {
|
|
||||||
"segments": [{"text": "Hello, World!"}]
|
|
||||||
}
|
|
||||||
|
|
||||||
# Mock diarization pipeline to raise an exception
|
|
||||||
mock_diarization.side_effect = Exception("Diarization failed")
|
|
||||||
|
|
||||||
stt = WhisperX(video_url)
|
|
||||||
transcription = stt.transcribe_youtube_video()
|
|
||||||
|
|
||||||
assert transcription == "Diarization failed"
|
|
||||||
|
|
||||||
|
|
||||||
# Add more tests for other scenarios and edge cases as needed.
|
|
||||||
Loading…
Reference in new issue