[OMNIMODALAGENT Fix]

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)
 ```

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
-

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",
                 )
-            )
+            )

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
-    

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"
+    assert instance.device == "cuda"