Kye
1 year ago
parent
f39d722f2a
commit
fd58cfa2a1
@ -0,0 +1,96 @@
|
|||||||
|
import time
|
||||||
|
import os
|
||||||
|
|
||||||
|
import pygame
|
||||||
|
import speech_recognition as sr
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from playsound import playsound
|
||||||
|
|
||||||
|
from swarms import OpenAIChat, OpenAITTS
|
||||||
|
|
||||||
|
# Load the environment variables
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
# Get the API key from the environment
|
||||||
|
openai_api_key = os.environ.get("OPENAI_API_KEY")
|
||||||
|
|
||||||
|
# Initialize the language model
|
||||||
|
llm = OpenAIChat(
|
||||||
|
openai_api_key=openai_api_key,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize the text-to-speech model
|
||||||
|
tts = OpenAITTS(
|
||||||
|
model_name="tts-1-1106",
|
||||||
|
voice="onyx",
|
||||||
|
openai_api_key=openai_api_key,
|
||||||
|
saved_filepath="runs/tts_speech.wav",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize the speech recognition model
|
||||||
|
r = sr.Recognizer()
|
||||||
|
|
||||||
|
|
||||||
|
def play_audio(file_path):
|
||||||
|
# Check if the file exists
|
||||||
|
if not os.path.isfile(file_path):
|
||||||
|
print(f"Audio file {file_path} not found.")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Initialize the mixer module
|
||||||
|
pygame.mixer.init()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Load the mp3 file
|
||||||
|
pygame.mixer.music.load(file_path)
|
||||||
|
|
||||||
|
# Play the mp3 file
|
||||||
|
pygame.mixer.music.play()
|
||||||
|
|
||||||
|
# Wait for the audio to finish playing
|
||||||
|
while pygame.mixer.music.get_busy():
|
||||||
|
pygame.time.Clock().tick(10)
|
||||||
|
except pygame.error as e:
|
||||||
|
print(f"Couldn't play {file_path}: {e}")
|
||||||
|
finally:
|
||||||
|
# Stop the mixer module and free resources
|
||||||
|
pygame.mixer.quit()
|
||||||
|
|
||||||
|
while True:
|
||||||
|
# Listen for user speech
|
||||||
|
with sr.Microphone() as source:
|
||||||
|
print("Listening...")
|
||||||
|
audio = r.listen(source)
|
||||||
|
|
||||||
|
# Convert speech to text
|
||||||
|
try:
|
||||||
|
print("Recognizing...")
|
||||||
|
task = r.recognize_google(audio)
|
||||||
|
print(f"User said: {task}")
|
||||||
|
except sr.UnknownValueError:
|
||||||
|
print("Could not understand audio")
|
||||||
|
continue
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error: {e}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
|
||||||
|
# Run the Gemini model on the task
|
||||||
|
print("Running GPT4 model...")
|
||||||
|
out = llm(task)
|
||||||
|
print(f"Gemini output: {out}")
|
||||||
|
|
||||||
|
# Convert the Gemini output to speech
|
||||||
|
print("Running text-to-speech model...")
|
||||||
|
out = tts.run_and_save(out)
|
||||||
|
print(f"Text-to-speech output: {out}")
|
||||||
|
|
||||||
|
# Ask the user if they want to play the audio
|
||||||
|
# play_audio = input("Do you want to play the audio? (yes/no): ")
|
||||||
|
# if play_audio.lower() == "yes":
|
||||||
|
# Initialize the mixer module
|
||||||
|
# Play the audio file
|
||||||
|
|
||||||
|
time.sleep(5)
|
||||||
|
|
||||||
|
playsound('runs/tts_speech.wav')
|
||||||
@ -1,336 +0,0 @@
|
|||||||
import os
|
|
||||||
import tempfile
|
|
||||||
from functools import wraps
|
|
||||||
from unittest.mock import AsyncMock, MagicMock, patch
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import pytest
|
|
||||||
import torch
|
|
||||||
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
|
|
||||||
|
|
||||||
from swarms.models.distilled_whisperx import (
|
|
||||||
DistilWhisperModel,
|
|
||||||
async_retry,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def distil_whisper_model():
|
|
||||||
return DistilWhisperModel()
|
|
||||||
|
|
||||||
|
|
||||||
def create_audio_file(
|
|
||||||
data: np.ndarray, sample_rate: int, file_path: str
|
|
||||||
):
|
|
||||||
data.tofile(file_path)
|
|
||||||
return file_path
|
|
||||||
|
|
||||||
|
|
||||||
def test_initialization(distil_whisper_model):
|
|
||||||
assert isinstance(distil_whisper_model, DistilWhisperModel)
|
|
||||||
assert isinstance(distil_whisper_model.model, torch.nn.Module)
|
|
||||||
assert isinstance(distil_whisper_model.processor, torch.nn.Module)
|
|
||||||
assert distil_whisper_model.device in ["cpu", "cuda:0"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_transcribe_audio_file(distil_whisper_model):
|
|
||||||
test_data = np.random.rand(
|
|
||||||
16000
|
|
||||||
) # Simulated audio data (1 second)
|
|
||||||
with tempfile.NamedTemporaryFile(
|
|
||||||
suffix=".wav", delete=False
|
|
||||||
) as audio_file:
|
|
||||||
audio_file_path = create_audio_file(
|
|
||||||
test_data, 16000, audio_file.name
|
|
||||||
)
|
|
||||||
transcription = distil_whisper_model.transcribe(
|
|
||||||
audio_file_path
|
|
||||||
)
|
|
||||||
os.remove(audio_file_path)
|
|
||||||
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
assert transcription.strip() != ""
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_transcribe_audio_file(distil_whisper_model):
|
|
||||||
test_data = np.random.rand(
|
|
||||||
16000
|
|
||||||
) # Simulated audio data (1 second)
|
|
||||||
with tempfile.NamedTemporaryFile(
|
|
||||||
suffix=".wav", delete=False
|
|
||||||
) as audio_file:
|
|
||||||
audio_file_path = create_audio_file(
|
|
||||||
test_data, 16000, audio_file.name
|
|
||||||
)
|
|
||||||
transcription = await distil_whisper_model.async_transcribe(
|
|
||||||
audio_file_path
|
|
||||||
)
|
|
||||||
os.remove(audio_file_path)
|
|
||||||
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
assert transcription.strip() != ""
|
|
||||||
|
|
||||||
|
|
||||||
def test_transcribe_audio_data(distil_whisper_model):
|
|
||||||
test_data = np.random.rand(
|
|
||||||
16000
|
|
||||||
) # Simulated audio data (1 second)
|
|
||||||
transcription = distil_whisper_model.transcribe(
|
|
||||||
test_data.tobytes()
|
|
||||||
)
|
|
||||||
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
assert transcription.strip() != ""
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_transcribe_audio_data(distil_whisper_model):
|
|
||||||
test_data = np.random.rand(
|
|
||||||
16000
|
|
||||||
) # Simulated audio data (1 second)
|
|
||||||
transcription = await distil_whisper_model.async_transcribe(
|
|
||||||
test_data.tobytes()
|
|
||||||
)
|
|
||||||
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
assert transcription.strip() != ""
|
|
||||||
|
|
||||||
|
|
||||||
def test_real_time_transcribe(distil_whisper_model, capsys):
|
|
||||||
test_data = np.random.rand(
|
|
||||||
16000 * 5
|
|
||||||
) # Simulated audio data (5 seconds)
|
|
||||||
with tempfile.NamedTemporaryFile(
|
|
||||||
suffix=".wav", delete=False
|
|
||||||
) as audio_file:
|
|
||||||
audio_file_path = create_audio_file(
|
|
||||||
test_data, 16000, audio_file.name
|
|
||||||
)
|
|
||||||
|
|
||||||
distil_whisper_model.real_time_transcribe(
|
|
||||||
audio_file_path, chunk_duration=1
|
|
||||||
)
|
|
||||||
|
|
||||||
os.remove(audio_file_path)
|
|
||||||
|
|
||||||
captured = capsys.readouterr()
|
|
||||||
assert "Starting real-time transcription..." in captured.out
|
|
||||||
assert "Chunk" in captured.out
|
|
||||||
|
|
||||||
|
|
||||||
def test_real_time_transcribe_audio_file_not_found(
|
|
||||||
distil_whisper_model, capsys
|
|
||||||
):
|
|
||||||
audio_file_path = "non_existent_audio.wav"
|
|
||||||
distil_whisper_model.real_time_transcribe(
|
|
||||||
audio_file_path, chunk_duration=1
|
|
||||||
)
|
|
||||||
|
|
||||||
captured = capsys.readouterr()
|
|
||||||
assert "The audio file was not found." in captured.out
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def mock_async_retry():
|
|
||||||
def _mock_async_retry(
|
|
||||||
retries=3, exceptions=(Exception,), delay=1
|
|
||||||
):
|
|
||||||
def decorator(func):
|
|
||||||
@wraps(func)
|
|
||||||
async def wrapper(*args, **kwargs):
|
|
||||||
return await func(*args, **kwargs)
|
|
||||||
|
|
||||||
return wrapper
|
|
||||||
|
|
||||||
return decorator
|
|
||||||
|
|
||||||
with patch(
|
|
||||||
"distil_whisper_model.async_retry", new=_mock_async_retry()
|
|
||||||
):
|
|
||||||
yield
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_retry_decorator_success():
|
|
||||||
async def mock_async_function():
|
|
||||||
return "Success"
|
|
||||||
|
|
||||||
decorated_function = async_retry()(mock_async_function)
|
|
||||||
result = await decorated_function()
|
|
||||||
assert result == "Success"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_retry_decorator_failure():
|
|
||||||
async def mock_async_function():
|
|
||||||
raise Exception("Error")
|
|
||||||
|
|
||||||
decorated_function = async_retry()(mock_async_function)
|
|
||||||
with pytest.raises(Exception, match="Error"):
|
|
||||||
await decorated_function()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_retry_decorator_multiple_attempts():
|
|
||||||
async def mock_async_function():
|
|
||||||
if mock_async_function.attempts == 0:
|
|
||||||
mock_async_function.attempts += 1
|
|
||||||
raise Exception("Error")
|
|
||||||
else:
|
|
||||||
return "Success"
|
|
||||||
|
|
||||||
mock_async_function.attempts = 0
|
|
||||||
decorated_function = async_retry(max_retries=2)(
|
|
||||||
mock_async_function
|
|
||||||
)
|
|
||||||
result = await decorated_function()
|
|
||||||
assert result == "Success"
|
|
||||||
|
|
||||||
|
|
||||||
def test_create_audio_file():
|
|
||||||
test_data = np.random.rand(
|
|
||||||
16000
|
|
||||||
) # Simulated audio data (1 second)
|
|
||||||
sample_rate = 16000
|
|
||||||
with tempfile.NamedTemporaryFile(
|
|
||||||
suffix=".wav", delete=False
|
|
||||||
) as audio_file:
|
|
||||||
audio_file_path = create_audio_file(
|
|
||||||
test_data, sample_rate, audio_file.name
|
|
||||||
)
|
|
||||||
|
|
||||||
assert os.path.exists(audio_file_path)
|
|
||||||
os.remove(audio_file_path)
|
|
||||||
|
|
||||||
|
|
||||||
# test_distilled_whisperx.py
|
|
||||||
|
|
||||||
|
|
||||||
# Fixtures for setting up model, processor, and audio files
|
|
||||||
@pytest.fixture(scope="module")
|
|
||||||
def model_id():
|
|
||||||
return "distil-whisper/distil-large-v2"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="module")
|
|
||||||
def whisper_model(model_id):
|
|
||||||
return DistilWhisperModel(model_id)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="session")
|
|
||||||
def audio_file_path(tmp_path_factory):
|
|
||||||
# You would create a small temporary MP3 file here for testing
|
|
||||||
# or use a public domain MP3 file's path
|
|
||||||
return "path/to/valid_audio.mp3"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="session")
|
|
||||||
def invalid_audio_file_path():
|
|
||||||
return "path/to/invalid_audio.mp3"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="session")
|
|
||||||
def audio_dict():
|
|
||||||
# This should represent a valid audio dictionary as expected by the model
|
|
||||||
return {"array": torch.randn(1, 16000), "sampling_rate": 16000}
|
|
||||||
|
|
||||||
|
|
||||||
# Test initialization
|
|
||||||
def test_initialization(whisper_model):
|
|
||||||
assert whisper_model.model is not None
|
|
||||||
assert whisper_model.processor is not None
|
|
||||||
|
|
||||||
|
|
||||||
# Test successful transcription with file path
|
|
||||||
def test_transcribe_with_file_path(whisper_model, audio_file_path):
|
|
||||||
transcription = whisper_model.transcribe(audio_file_path)
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
|
|
||||||
|
|
||||||
# Test successful transcription with audio dict
|
|
||||||
def test_transcribe_with_audio_dict(whisper_model, audio_dict):
|
|
||||||
transcription = whisper_model.transcribe(audio_dict)
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
|
|
||||||
|
|
||||||
# Test for file not found error
|
|
||||||
def test_file_not_found(whisper_model, invalid_audio_file_path):
|
|
||||||
with pytest.raises(Exception):
|
|
||||||
whisper_model.transcribe(invalid_audio_file_path)
|
|
||||||
|
|
||||||
|
|
||||||
# Asynchronous tests
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_transcription_success(
|
|
||||||
whisper_model, audio_file_path
|
|
||||||
):
|
|
||||||
transcription = await whisper_model.async_transcribe(
|
|
||||||
audio_file_path
|
|
||||||
)
|
|
||||||
assert isinstance(transcription, str)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_transcription_failure(
|
|
||||||
whisper_model, invalid_audio_file_path
|
|
||||||
):
|
|
||||||
with pytest.raises(Exception):
|
|
||||||
await whisper_model.async_transcribe(invalid_audio_file_path)
|
|
||||||
|
|
||||||
|
|
||||||
# Testing real-time transcription simulation
|
|
||||||
def test_real_time_transcription(
|
|
||||||
whisper_model, audio_file_path, capsys
|
|
||||||
):
|
|
||||||
whisper_model.real_time_transcribe(
|
|
||||||
audio_file_path, chunk_duration=1
|
|
||||||
)
|
|
||||||
captured = capsys.readouterr()
|
|
||||||
assert "Starting real-time transcription..." in captured.out
|
|
||||||
|
|
||||||
|
|
||||||
# Testing retry decorator for asynchronous function
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_retry():
|
|
||||||
@async_retry(max_retries=2, exceptions=(ValueError,), delay=0)
|
|
||||||
async def failing_func():
|
|
||||||
raise ValueError("Test")
|
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
|
||||||
await failing_func()
|
|
||||||
|
|
||||||
|
|
||||||
# Mocking the actual model to avoid GPU/CPU intensive operations during test
|
|
||||||
@pytest.fixture
|
|
||||||
def mocked_model(monkeypatch):
|
|
||||||
model_mock = AsyncMock(AutoModelForSpeechSeq2Seq)
|
|
||||||
processor_mock = MagicMock(AutoProcessor)
|
|
||||||
monkeypatch.setattr(
|
|
||||||
"swarms.models.distilled_whisperx.AutoModelForSpeechSeq2Seq.from_pretrained",
|
|
||||||
model_mock,
|
|
||||||
)
|
|
||||||
monkeypatch.setattr(
|
|
||||||
"swarms.models.distilled_whisperx.AutoProcessor.from_pretrained",
|
|
||||||
processor_mock,
|
|
||||||
)
|
|
||||||
return model_mock, processor_mock
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_transcribe_with_mocked_model(
|
|
||||||
mocked_model, audio_file_path
|
|
||||||
):
|
|
||||||
model_mock, processor_mock = mocked_model
|
|
||||||
# Set up what the mock should return when it's called
|
|
||||||
model_mock.return_value.generate.return_value = torch.tensor(
|
|
||||||
[[0]]
|
|
||||||
)
|
|
||||||
processor_mock.return_value.batch_decode.return_value = [
|
|
||||||
"mocked transcription"
|
|
||||||
]
|
|
||||||
model_wrapper = DistilWhisperModel()
|
|
||||||
transcription = await model_wrapper.async_transcribe(
|
|
||||||
audio_file_path
|
|
||||||
)
|
|
||||||
assert transcription == "mocked transcription"
|
|
||||||
Loading…
Reference in new issue