diff --git a/playground/agents/tool_agent.py b/playground/agents/tool_agent.py
index 93e07ff3..0a95f42c 100644
--- a/playground/agents/tool_agent.py
+++ b/playground/agents/tool_agent.py
@@ -4,7 +4,9 @@ from swarms import ToolAgent
 
 # Load the pre-trained model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
-    "databricks/dolly-v2-12b"
+    "databricks/dolly-v2-12b",
+    load_in_4bit=True,
+    device_map="auto",
 )
 tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
 
@@ -26,11 +28,15 @@ task = (
 
 # Create an instance of the ToolAgent class
 agent = ToolAgent(
-    model=model, tokenizer=tokenizer, json_schema=json_schema
+    name="dolly-function-agent",
+    description="Ana gent to create a child data",
+    model=model,
+    tokenizer=tokenizer,
+    json_schema=json_schema,
 )
 
 # Run the agent to generate the person's information
 generated_data = agent.run(task)
 
 # Print the generated data
-print(generated_data)
+print(f"Generated data: {generated_data}")
diff --git a/playground/models/kosmos2_example.py b/playground/models/kosmos2_example.py
deleted file mode 100644
index 6fc4df02..00000000
--- a/playground/models/kosmos2_example.py
+++ /dev/null
@@ -1,10 +0,0 @@
-from swarms.models.kosmos2 import Kosmos2
-from PIL import Image
-
-
-model = Kosmos2.initialize()
-
-image = Image.open("images/swarms.jpg")
-
-detections = model(image)
-print(detections)
diff --git a/playground/models/kosmos_two_example.py b/playground/models/kosmos_two_example.py
deleted file mode 100644
index 8bf583cd..00000000
--- a/playground/models/kosmos_two_example.py
+++ /dev/null
@@ -1,11 +0,0 @@
-from swarms.models.kosmos_two import Kosmos
-
-# Initialize Kosmos
-kosmos = Kosmos()
-
-# Perform multimodal grounding
-out = kosmos.multimodal_grounding(
-    "Find the red apple in the image.", "images/swarms.jpeg"
-)
-
-print(out)
diff --git a/pyproject.toml b/pyproject.toml
index c9f0c28d..f4971b47 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "4.0.1"
+version = "4.0.2"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
diff --git a/swarms/agents/tool_agent.py b/swarms/agents/tool_agent.py
index d9359471..8e6adf9d 100644
--- a/swarms/agents/tool_agent.py
+++ b/swarms/agents/tool_agent.py
@@ -61,11 +61,11 @@ class ToolAgent(AbstractLLM):
 
     def __init__(
         self,
-        name: str,
-        description: str,
-        model: Any,
-        tokenizer: Any,
-        json_schema: Any,
+        name: str = "Function Calling Agent",
+        description: str = "Generates a function based on the input json schema and the task",
+        model: Any = None,
+        tokenizer: Any = None,
+        json_schema: Any = None,
         max_number_tokens: int = 500,
         *args,
         **kwargs,