diff --git a/docs/swarms/agents/omni_agent.md b/docs/swarms/agents/omni_agent.md index 93ae18dc..a80da0e6 100644 --- a/docs/swarms/agents/omni_agent.md +++ b/docs/swarms/agents/omni_agent.md @@ -39,9 +39,26 @@ For streaming mode, this function yields the response token by token, ensuring a ## Examples & Use Cases Initialize the `OmniModalAgent` and communicate with it: ```python -from swarms import OmniModalAgent, OpenAIChat -llm_instance = OpenAIChat() -agent = OmniModalAgent(llm_instance) +from swarms.agents.omni_modal_agent import OmniModalAgent, OpenAIChat +from swarms.models import OpenAIChat +from dotenv import load_dotenv +import os + +# Load the environment variables +load_dotenv() + +# Get the API key from the environment +api_key = os.environ.get("OPENAI_API_KEY") + +# Initialize the language model +llm = OpenAIChat( + temperature=0.5, + model_name="gpt-4", + openai_api_key=api_key, +) + + +agent = OmniModalAgent(llm) response = agent.run("Translate 'Hello' to French.") print(response) ``` diff --git a/swarms/models/base_multimodal_model.py b/swarms/models/base_multimodal_model.py index 28b21d64..2f6110d6 100644 --- a/swarms/models/base_multimodal_model.py +++ b/swarms/models/base_multimodal_model.py @@ -311,4 +311,3 @@ class BaseMultiModalModel: def set_max_length(self, max_length): """Set max_length""" self.max_length = max_length - diff --git a/swarms/models/fuyu.py b/swarms/models/fuyu.py index 867ce5c0..862ed882 100644 --- a/swarms/models/fuyu.py +++ b/swarms/models/fuyu.py @@ -1,4 +1,3 @@ - from PIL import Image from termcolor import colored from transformers import ( @@ -15,7 +14,7 @@ class Fuyu(BaseMultiModalModel): """ Fuyu model by Adept - + Args: BaseMultiModalModel (BaseMultiModalModel): [description] model_name (str, optional): [description]. Defaults to "adept/fuyu-8b". @@ -23,9 +22,9 @@ class Fuyu(BaseMultiModalModel): max_new_tokens (int, optional): [description]. Defaults to 500. *args: [description] **kwargs: [description] - - - + + + Examples: >>> from swarms.models import Fuyu >>> model = Fuyu() @@ -45,9 +44,7 @@ class Fuyu(BaseMultiModalModel): self.device_map = device_map self.max_new_tokens = max_new_tokens - self.tokenizer = AutoTokenizer.from_pretrained( - model_name - ) + self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.image_processor = FuyuImageProcessor() self.processor = FuyuProcessor( image_processor=self.image_processor, @@ -111,4 +108,4 @@ class Fuyu(BaseMultiModalModel): ), "red", ) - ) \ No newline at end of file + ) diff --git a/tests/models/test_fuyu.py b/tests/models/test_fuyu.py index 6c653f28..e76e11bb 100644 --- a/tests/models/test_fuyu.py +++ b/tests/models/test_fuyu.py @@ -1,4 +1,3 @@ - from unittest.mock import patch import pytest @@ -127,7 +126,6 @@ def test_processor_initialization(fuyu_instance): assert isinstance(fuyu_instance.processor, FuyuProcessor) - # Test `get_img` method with a valid image path def test_get_img_valid_path(fuyu_instance): with patch("PIL.Image.open") as mock_open: @@ -135,6 +133,7 @@ def test_get_img_valid_path(fuyu_instance): result = fuyu_instance.get_img("valid/path/to/image.png") assert result == "Test image" + # Test `get_img` method with an invalid image path def test_get_img_invalid_path(fuyu_instance): with patch("PIL.Image.open") as mock_open: @@ -142,39 +141,55 @@ def test_get_img_invalid_path(fuyu_instance): with pytest.raises(FileNotFoundError): fuyu_instance.get_img("invalid/path/to/image.png") + # Test `run` method with valid inputs def test_run_valid_inputs(fuyu_instance): - with patch.object(fuyu_instance, "get_img") as mock_get_img, \ - patch.object(fuyu_instance, "processor") as mock_processor, \ - patch.object(fuyu_instance, "model") as mock_model: + with patch.object( + fuyu_instance, "get_img" + ) as mock_get_img, patch.object( + fuyu_instance, "processor" + ) as mock_processor, patch.object( + fuyu_instance, "model" + ) as mock_model: mock_get_img.return_value = "Test image" - mock_processor.return_value = {"input_ids": torch.tensor([1, 2, 3])} + mock_processor.return_value = { + "input_ids": torch.tensor([1, 2, 3]) + } mock_model.generate.return_value = torch.tensor([1, 2, 3]) mock_processor.batch_decode.return_value = ["Test text"] - result = fuyu_instance.run("Hello, world!", "valid/path/to/image.png") + result = fuyu_instance.run( + "Hello, world!", "valid/path/to/image.png" + ) assert result == ["Test text"] + # Test `run` method with invalid text input def test_run_invalid_text_input(fuyu_instance): with pytest.raises(Exception): fuyu_instance.run(None, "valid/path/to/image.png") + # Test `run` method with empty text input def test_run_empty_text_input(fuyu_instance): with pytest.raises(Exception): fuyu_instance.run("", "valid/path/to/image.png") + # Test `run` method with very long text input def test_run_very_long_text_input(fuyu_instance): with pytest.raises(Exception): fuyu_instance.run("A" * 10000, "valid/path/to/image.png") + # Test `run` method with invalid image path def test_run_invalid_image_path(fuyu_instance): with patch.object(fuyu_instance, "get_img") as mock_get_img: mock_get_img.side_effect = FileNotFoundError with pytest.raises(FileNotFoundError): - fuyu_instance.run("Hello, world!", "invalid/path/to/image.png") + fuyu_instance.run( + "Hello, world!", "invalid/path/to/image.png" + ) + # Test `__init__` method with default parameters def test_init_default_parameters(): @@ -183,10 +198,10 @@ def test_init_default_parameters(): assert fuyu_instance.device_map == "auto" assert fuyu_instance.max_new_tokens == 500 + # Test `__init__` method with custom parameters def test_init_custom_parameters(): fuyu_instance = Fuyu("custom/path", "cpu", 1000) assert fuyu_instance.pretrained_path == "custom/path" assert fuyu_instance.device_map == "cpu" assert fuyu_instance.max_new_tokens == 1000 - diff --git a/tests/models/test_idefics.py b/tests/models/test_idefics.py index 3a0a8cdc..25a8dd5b 100644 --- a/tests/models/test_idefics.py +++ b/tests/models/test_idefics.py @@ -157,16 +157,18 @@ def test_run_batched_mode_false(idefics_instance): assert result == ["Test"] + # Test `run` method with an exception def test_run_with_exception(idefics_instance): task = "User: Test" with patch.object( idefics_instance, "processor" ) as mock_processor: - mock_processor.side_effect = Exception('Test exception') + mock_processor.side_effect = Exception("Test exception") with pytest.raises(Exception): idefics_instance.run(task) + # Test `set_model_name` method def test_set_model_name(idefics_instance): new_model_name = "new_model_name" @@ -182,6 +184,7 @@ def test_set_model_name(idefics_instance): new_model_name, torch_dtype=torch.bfloat16 ) + # Test `__init__` method with device set to None def test_init_device_none(): with patch( @@ -191,6 +194,7 @@ def test_init_device_none(): instance = Idefics(device=None) assert instance.device == "cpu" + # Test `__init__` method with device set to "cuda" def test_init_device_cuda(): with patch( @@ -198,4 +202,4 @@ def test_init_device_cuda(): return_value=True, ): instance = Idefics(device="cuda") - assert instance.device == "cuda" \ No newline at end of file + assert instance.device == "cuda"