SqueezeAILab / TinyAgent

[EMNLP 2024 Demo] TinyAgent: Function Calling at the Edge!
https://arxiv.org/abs/2409.00608
MIT License
311 stars 49 forks source link

TinyAgent: Function Calling at the Edge

Arxiv GitHub license

Get the desktop app‎ ‎ |‎ ‎ Read the blog post‎ ‎ |‎ ‎ Read the paper

Thumbnail

TinyAgent aims to enable complex reasoning and function calling capabilities in Small Language Models (SLMs) that can be deployed securely and privately at the edge. Traditional Large Language Models (LLMs) like GPT-4 and Gemini-1.5, while powerful, are often too large and resource-intensive for edge deployment, posing challenges in terms of privacy, connectivity, and latency. TinyAgent addresses these challenges by training specialized SLMs with high-quality, curated data, and focusing on function calling with LLMCompiler. As a driving application, TinyAgent can interact with various MacOS applications, assisting users with day-to-day tasks such as composing emails, managing contacts, scheduling calendar events, and organizing Zoom meetings.

Demo

TinyAgent Demo

What can TinyAgent do?

TinyAgent is equipped with 16 different functions that can interact with different applications on Mac, which includes:

📧 Mail

📇 Contacts

📨 SMS

📅 Calendar

🗺️ Maps

🗒️ Notes

🗂️ File Management

⏰ Reminders

🎥 Zoom Meetings

💬 Custom Instructions

You can choose to enable/disable certain apps by going to the Preferences window.

🤖 Sub-Agents

Depending on the task simplicity, TinyAgent orchestrates the execution of different more specialized or smaller LMs. TinyAgent currently can operate LMs that can summarize a PDF, write emails, or take notes.

See the Customization section to see how to add your own sub-agentx.

🛠️ ToolRAG

When faced with challenging tasks, SLM agents require appropriate tools and in-context examples to guide them. If the model sees irrelevant examples, it can hallucinate. Likewise, if the model sees the descriptions of the tools that it doesn’t need, it usually gets confused, and these tools take up unnecessary prompt space. To tackle this, TinyAgent uses ToolRAG to retrieve the best tools and examples suited for a given query. This process has minimal latency and increases the accuracy of TinyAgent substantially. Please take a look at our blog post and our ToolRAG model for more details.

You need to first install our ToolRAG model from Hugging Face and enable it from the TinyAgent settings to use it.

🎙️ Whisper

TinyAgent also accepts voice commands through both the OpenAI Whisper API and local whisper.cpp deployment. For whisper.cpp, you need to setup the local whisper server and provide the server port number in the TinyAgent settings.

Providers

You can use with your OpenAI key, Azure deployments, or even your own local models!

OpenAI

You need to provide OpenAI API Key and the models you want to use in the 'Preferences' window.

Azure Deployments

You need to provide your deployment name and the endpoints for the main agent/sub-agents/embedding model as well as the context length of the agent models in the 'Preferences' window.

Local Models

You can plug-and-play every part of TinyAgent with your local models! TinyAgent can use an OpenAI-compatible server to run models locally. There are several options you can take:

All TinyAgent needs is the port numbers that you are serving your model at and its context length.

Fine-tuned TinyAgents

We also provide our own fine-tuned open source models, TinyAgent-1.1B and TinyAgent-7B! We curated a dataset of 40000 real-life use cases for TinyAgent and fine-tuned two small open-source language models on this dataset with LoRA. After fine-tuning and using ToolRAG, both TinyAgent-1.1B and TinyAgent-7B exceed the performance of GPT-4-turbo. Check out our for the specifics of dataset generation, evaluation, and fine-tuning.

Model Success Rate
GPT-3.5-turbo 65.04%
GPT-4-turbo 79.08%
TinyLLama-1.1B-32K-Instruct 12.71%
WizardLM-2-7b 41.25%
TinyAgent-1.1B + ToolRAG / [hf] [gguf] 80.06%
TinyAgent-7B + ToolRAG / [hf] [gguf] 84.95%

Customization

You can customize your TinyAgent by going to ~/Library/Application Support/TinyAgent/tinyagent-llmcompiler directory and changing the code yourself.

Using TinyAgent programmatically

You can use TinyAgent programmatically by just passing in a config file.

from src.tiny_agent.tiny_agent import TinyAgent
from src.tiny_agent.config import get_tiny_agent_config

config_path = "..."
tiny_agent_config = get_tiny_agent_config(config_path=config_path)
tiny_agent = TinyAgent(tiny_agent_config)

await TinyAgent.arun(query="Create a meeting with Sid and Lutfi for tomorrow 2pm to discuss the meeting notes.")

Adding your own tools

  1. Navigate to src/tiny_agent/models.py and add your tool;s name to TinyAgentToolName(Enum)
class TinyAgentToolName(Enum):
...
CUSTOM_TOOL_NAME = "custom_tool"
  1. Navigate to src/tiny_agent/tiny_agent_tools.py and define your tool.
def get_custom_tool(...) -> list[Tool]:
async def tool_coroutine(...) -> str: # Needs to return a string
...
return ...

    custom_tool = Tool(
        name=TinyAgentToolName.CUSTOM_TOOL_NAME,
        func=tool_coroutine,
        description=(
            f"{TinyAgentToolName.CUSTOM_TOOL_NAME.value}(...) -> str\n"
            "<The description of the tool>"
        )
    )

    return [custom_tool]
  1. Add your tools to the get_tiny_agent_tools function.
def get_tiny_agent_tools(...):
...
tools += get_custom_tools(...)
...

Note: Adding your own tools only works for GPT models since our open-source models and ToolRAG were only fine-tuned on the original TinyAgent toolset.

Adding your own sub-agents

You can also add your own custom subagents. To do so, please follow these steps:

  1. All sub-agents inherit from SubAgent class in src/tiny_agent/sub_agents/sub_agent.py. Your custom agent should inherit from this abstract class and define the __call__ method.
class CustomSubAgent(SubAgent):

    async def __call__(self, ...) -> str:
        ...
        return response
  1. Add your custom agent to TinyAgent in src/tiny_agent/tiny_agent.py
from src.tiny_agent.sub_agents.custom_agent import CustomAgent

class TinyAgent:
...
custom_agent: CustomAgent

    def __init__(...):
        ...
        self.custom_agent = CustomAgent(
            sub_agent_llm, config.sub_agent_config, config.custom_instructions
        )
        ...
  1. After defining your custom agent and adding it to TinyAgent, you should create a tool that calls this agent. Please refer to the Adding your own tools section to see how to do so

Citation

We would appreciate it if you could please cite our paper if you found TinyAgent useful for your work:

@misc{erdogan2024tinyagentfunctioncallingedge,
      title={TinyAgent: Function Calling at the Edge},
      author={Lutfi Eren Erdogan and Nicholas Lee and Siddharth Jha and Sehoon Kim and Ryan Tabrizi and Suhong Moon and Coleman Hooper and Gopala Anumanchipalli and Kurt Keutzer and Amir Gholami},
      year={2024},
      eprint={2409.00608},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2409.00608},
}