tests for yi, stable diffusion, timm models, etc

Former-commit-id: dfea671

swarms/models/autotemp.py

@@ -1,6 +1,6 @@
 import re
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from swarms.models.auto_temp import OpenAIChat
+from swarms.models.openai_models import OpenAIChat
 
 
 class AutoTempAgent:

swarms/models/simple_ada.py

@@ -2,6 +2,7 @@
 
 client = OpenAI()
 
+
 def get_ada_embeddings(text: str, model: str = "text-embedding-ada-002"):
  """
  Simple function to get embeddings from ada

tests/models/bioclip.py

@@ -0,0 +1,161 @@
+# Import necessary modules and define fixtures if needed
+import os
+import pytest
+import torch
+from PIL import Image
+from swarms.models.bioclip import BioClip
+
+
+# Define fixtures if needed
+@pytest.fixture
+def sample_image_path():
+ return "path_to_sample_image.jpg"
+
+
+@pytest.fixture
+def clip_instance():
+ return BioClip("microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
+
+
+# Basic tests for the BioClip class
+def test_clip_initialization(clip_instance):
+ assert isinstance(clip_instance.model, torch.nn.Module)
+ assert hasattr(clip_instance, "model_path")
+ assert hasattr(clip_instance, "preprocess_train")
+ assert hasattr(clip_instance, "preprocess_val")
+ assert hasattr(clip_instance, "tokenizer")
+ assert hasattr(clip_instance, "device")
+
+
+def test_clip_call_method(clip_instance, sample_image_path):
+ labels = [
+ "adenocarcinoma histopathology",
+ "brain MRI",
+ "covid line chart",
+ "squamous cell carcinoma histopathology",
+ "immunohistochemistry histopathology",
+ "bone X-ray",
+ "chest X-ray",
+ "pie chart",
+ "hematoxylin and eosin histopathology",
+ ]
+ result = clip_instance(sample_image_path, labels)
+ assert isinstance(result, dict)
+ assert len(result) == len(labels)
+
+
+def test_clip_plot_image_with_metadata(clip_instance, sample_image_path):
+ metadata = {
+ "filename": "sample_image.jpg",
+ "top_probs": {"label1": 0.75, "label2": 0.65},
+ }
+ clip_instance.plot_image_with_metadata(sample_image_path, metadata)
+
+
+# More test cases can be added to cover additional functionality and edge cases
+
+
+# Parameterized tests for different image and label combinations
+@pytest.mark.parametrize(
+ "image_path, labels",
+ [
+ ("image1.jpg", ["label1", "label2"]),
+ ("image2.jpg", ["label3", "label4"]),
+ # Add more image and label combinations
+ ],
+)
+def test_clip_parameterized_calls(clip_instance, image_path, labels):
+ result = clip_instance(image_path, labels)
+ assert isinstance(result, dict)
+ assert len(result) == len(labels)
+
+
+# Test image preprocessing
+def test_clip_image_preprocessing(clip_instance, sample_image_path):
+ image = Image.open(sample_image_path)
+ processed_image = clip_instance.preprocess_val(image)
+ assert isinstance(processed_image, torch.Tensor)
+
+
+# Test label tokenization
+def test_clip_label_tokenization(clip_instance):
+ labels = ["label1", "label2"]
+ tokenized_labels = clip_instance.tokenizer(labels)
+ assert isinstance(tokenized_labels, torch.Tensor)
+ assert tokenized_labels.shape[0] == len(labels)
+
+
+# More tests can be added to cover other methods and edge cases
+
+
+# End-to-end tests with actual images and labels
+def test_clip_end_to_end(clip_instance, sample_image_path):
+ labels = [
+ "adenocarcinoma histopathology",
+ "brain MRI",
+ "covid line chart",
+ "squamous cell carcinoma histopathology",
+ "immunohistochemistry histopathology",
+ "bone X-ray",
+ "chest X-ray",
+ "pie chart",
+ "hematoxylin and eosin histopathology",
+ ]
+ result = clip_instance(sample_image_path, labels)
+ assert isinstance(result, dict)
+ assert len(result) == len(labels)
+
+
+# Test label tokenization with long labels
+def test_clip_long_labels(clip_instance):
+ labels = ["label" + str(i) for i in range(100)]
+ tokenized_labels = clip_instance.tokenizer(labels)
+ assert isinstance(tokenized_labels, torch.Tensor)
+ assert tokenized_labels.shape[0] == len(labels)
+
+
+# Test handling of multiple image files
+def test_clip_multiple_images(clip_instance, sample_image_path):
+ labels = ["label1", "label2"]
+ image_paths = [sample_image_path, "image2.jpg"]
+ results = clip_instance(image_paths, labels)
+ assert isinstance(results, list)
+ assert len(results) == len(image_paths)
+ for result in results:
+ assert isinstance(result, dict)
+ assert len(result) == len(labels)
+
+
+# Test model inference performance
+def test_clip_inference_performance(clip_instance, sample_image_path, benchmark):
+ labels = [
+ "adenocarcinoma histopathology",
+ "brain MRI",
+ "covid line chart",
+ "squamous cell carcinoma histopathology",
+ "immunohistochemistry histopathology",
+ "bone X-ray",
+ "chest X-ray",
+ "pie chart",
+ "hematoxylin and eosin histopathology",
+ ]
+ result = benchmark(clip_instance, sample_image_path, labels)
+ assert isinstance(result, dict)
+ assert len(result) == len(labels)
+
+
+# Test different preprocessing pipelines
+def test_clip_preprocessing_pipelines(clip_instance, sample_image_path):
+ labels = ["label1", "label2"]
+ image = Image.open(sample_image_path)
+
+ # Test preprocessing for training
+ processed_image_train = clip_instance.preprocess_train(image)
+ assert isinstance(processed_image_train, torch.Tensor)
+
+ # Test preprocessing for validation
+ processed_image_val = clip_instance.preprocess_val(image)
+ assert isinstance(processed_image_val, torch.Tensor)
+
+
+# ...

tests/models/distill_whisper.py

@@ -1,13 +1,14 @@
 import os
 import tempfile
 from functools import wraps
-from unittest.mock import patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
 
-from swarms.models.distill_whisperx import DistilWhisperModel, async_retry
+from swarms.models.distilled_whisperx import DistilWhisperModel, async_retry
 
 
 @pytest.fixture
@@ -150,5 +151,114 @@ def test_create_audio_file():
  os.remove(audio_file_path)
 
 
-if __name__ == "__main__":
- pytest.main()
+# 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"