From 8b3b2fde1e5b0a63a7a7af3dc81a5314114db7c8 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Wed, 31 Jan 2024 21:18:33 -0500
Subject: [PATCH] [EXAMPLE][ToolAgent]

---
 README.md                       | 18 +++++++++++------
 playground/agents/tool_agent.py | 36 +++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100644 playground/agents/tool_agent.py

diff --git a/README.md b/README.md
index 52fc6a04..1b22bc10 100644
--- a/README.md
+++ b/README.md
@@ -83,32 +83,38 @@ ToolAgent is an agent that outputs JSON using any model from huggingface. It tak
 
 
 ```python
+# Import necessary libraries
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from swarms import ToolAgent
 
-
+# Load the pre-trained model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
 tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
 
+# Define a JSON schema for person's information
 json_schema = {
     "type": "object",
     "properties": {
         "name": {"type": "string"},
         "age": {"type": "number"},
         "is_student": {"type": "boolean"},
-        "courses": {
-            "type": "array",
-            "items": {"type": "string"}
-        }
-    }
+        "courses": {"type": "array", "items": {"type": "string"}},
+    },
 }
 
+# Define the task to generate a person's information
 task = "Generate a person's information based on the following schema:"
+
+# Create an instance of the ToolAgent class
 agent = ToolAgent(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)
 
+
 ```
 
 
diff --git a/playground/agents/tool_agent.py b/playground/agents/tool_agent.py
new file mode 100644
index 00000000..93e07ff3
--- /dev/null
+++ b/playground/agents/tool_agent.py
@@ -0,0 +1,36 @@
+# Import necessary libraries
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from swarms import ToolAgent
+
+# Load the pre-trained model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(
+    "databricks/dolly-v2-12b"
+)
+tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
+
+# Define a JSON schema for person's information
+json_schema = {
+    "type": "object",
+    "properties": {
+        "name": {"type": "string"},
+        "age": {"type": "number"},
+        "is_student": {"type": "boolean"},
+        "courses": {"type": "array", "items": {"type": "string"}},
+    },
+}
+
+# Define the task to generate a person's information
+task = (
+    "Generate a person's information based on the following schema:"
+)
+
+# Create an instance of the ToolAgent class
+agent = ToolAgent(
+    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)