From a86f5979af4cb66a7c104ed629abd1a77e6b62b4 Mon Sep 17 00:00:00 2001 From: Kye Date: Thu, 4 Apr 2024 19:09:36 -0400 Subject: [PATCH] [EXAMPLES][Jamba ++ Command R Tool Agent] --- playground/agents/command_r_tool_agent.py | 59 +++++++++++++++++++++++ playground/agents/jamba_tool_agent.py | 59 +++++++++++++++++++++++ pyproject.toml | 4 +- swarms/utils/__init__.py | 4 +- 4 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 playground/agents/command_r_tool_agent.py create mode 100644 playground/agents/jamba_tool_agent.py diff --git a/playground/agents/command_r_tool_agent.py b/playground/agents/command_r_tool_agent.py new file mode 100644 index 00000000..cd9644db --- /dev/null +++ b/playground/agents/command_r_tool_agent.py @@ -0,0 +1,59 @@ +from pydantic import BaseModel, Field +from transformers import AutoModelForCausalLM, AutoTokenizer + +from swarms import ToolAgent +from swarms.utils.json_utils import base_model_to_json + +# Model name +model_name = "CohereForAI/c4ai-command-r-v01-4bit" + +# Load the pre-trained model and tokenizer +model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map="auto", +) + +# Load the pre-trained model and tokenizer +tokenizer = AutoTokenizer.from_pretrained(model_name) + +# Initialize the schema for the person's information +class APIExampleRequestSchema(BaseModel): + endpoint: str = Field( + ..., description="The API endpoint for the example request" + ) + method: str = Field( + ..., description="The HTTP method for the example request" + ) + headers: dict = Field( + ..., description="The headers for the example request" + ) + body: dict = Field( + ..., description="The body of the example request" + ) + response: dict = Field( + ..., description="The expected response of the example request" + ) + +# Convert the schema to a JSON string +api_example_schema = base_model_to_json(APIExampleRequestSchema) +# Convert the schema to a JSON string + +# Define the task to generate a person's information +task = ( + "Generate an example API request using this code:\n" +) + +# Create an instance of the ToolAgent class +agent = ToolAgent( + name="Command R Tool Agent", + description="An agent that generates an API request using the Command R model.", + model=model, + tokenizer=tokenizer, + json_schema=api_example_schema, +) + +# Run the agent to generate the person's information +generated_data = agent.run(task) + +# Print the generated data +print(f"Generated data: {generated_data}") \ No newline at end of file diff --git a/playground/agents/jamba_tool_agent.py b/playground/agents/jamba_tool_agent.py new file mode 100644 index 00000000..1bc2666c --- /dev/null +++ b/playground/agents/jamba_tool_agent.py @@ -0,0 +1,59 @@ +from pydantic import BaseModel, Field +from transformers import AutoModelForCausalLM, AutoTokenizer + +from swarms import ToolAgent +from swarms.utils.json_utils import base_model_to_json + +# Model name +model_name = "ai21labs/Jamba-v0.1" + +# Load the pre-trained model and tokenizer +model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map="auto", +) + +# Load the pre-trained model and tokenizer +tokenizer = AutoTokenizer.from_pretrained(model_name) + +# Initialize the schema for the person's information +class APIExampleRequestSchema(BaseModel): + endpoint: str = Field( + ..., description="The API endpoint for the example request" + ) + method: str = Field( + ..., description="The HTTP method for the example request" + ) + headers: dict = Field( + ..., description="The headers for the example request" + ) + body: dict = Field( + ..., description="The body of the example request" + ) + response: dict = Field( + ..., description="The expected response of the example request" + ) + +# Convert the schema to a JSON string +api_example_schema = base_model_to_json(APIExampleRequestSchema) +# Convert the schema to a JSON string + +# Define the task to generate a person's information +task = ( + "Generate an example API request using this code:\n" +) + +# Create an instance of the ToolAgent class +agent = ToolAgent( + name="Command R Tool Agent", + description="An agent that generates an API request using the Command R model.", + model=model, + tokenizer=tokenizer, + json_schema=api_example_schema, +) + +# Run the agent to generate the person's information +generated_data = agent(task) + +# Print the generated data +print(f"Generated data: {generated_data}") \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index d1b8b199..238b317b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] name = "swarms" -version = "4.7.4" +version = "4.7.8" description = "Swarms - Pytorch" license = "MIT" authors = ["Kye Gomez "] @@ -44,6 +44,8 @@ Pillow = "10.2.0" rich = "13.5.2" psutil = "*" sentry-sdk = "*" +python-dotenv = "*" +accelerate = "0.28.0" [tool.poetry.dev-dependencies] black = "23.3.0" diff --git a/swarms/utils/__init__.py b/swarms/utils/__init__.py index 7a057975..945479fe 100644 --- a/swarms/utils/__init__.py +++ b/swarms/utils/__init__.py @@ -40,7 +40,7 @@ from swarms.utils.remove_json_whitespace import ( remove_whitespace_from_yaml, ) from swarms.utils.save_logs import parse_log_file -from swarms.utils.supervision_visualizer import MarkVisualizer +# from swarms.utils.supervision_visualizer import MarkVisualizer from swarms.utils.try_except_wrapper import try_except_wrapper from swarms.utils.yaml_output_parser import YamlOutputParser from swarms.utils.concurrent_utils import execute_concurrently @@ -77,8 +77,6 @@ __all__ = [ "remove_whitespace_from_json", "remove_whitespace_from_yaml", "parse_log_file", - "MarkVisualizer", - # "limit_tokens_from_string", "try_except_wrapper", "YamlOutputParser", "execute_concurrently",