Efficient and Scalable Estimation of Tool Representations in Vector Space

Efficient and scalable tool retrieval is critical for modern function calling applications. We propose novel approaches to the tool retrieval problem: (1) Tool2Vec: usage-driven tool embedding generation for tool retrieval, (2) ToolRefiner: a staged retrieval method that iteratively improves the quality of retrieved tools, and (3) MLC: framing tool retrieval as a multi-label classification problem. With these new methods, we achieve improvements of up to 27.28 in Recall@K on the ToolBench dataset. Furthermore, we introduce ToolBank, a set of domain-specific tool retrieval datasets to encourage further research. For more details, please check out our paper here.

Installation

Create a conda environment and install the dependencies

conda create --name ToolRAG python=3.10 -y
conda activate ToolRAG

Clone and install the dependencies

git clone https://github.com/SuhongMoon/ToolRAG.git
cd ToolRAG
pip install -e .
pip install -r requirements.txt

Download ToolBank Dataset

Install HuggingFace datasets package

pip install datasets

Load the dataset from HuggingFace

from datasets import load_dataset

tool_bank = load_dataset("squeeze-ai-lab/ToolBank")

The link to the dataset: here

Basic Runs

Generate Synthetic Data

Refer to toolrag/data_generation/README.md for more details here.

Fine Tuning Embedding Models

Refer to toolrag/query2query/README.md for more details here.

Generate Tool2Vec Embeddings

Refer to toolrag/tool2vec/README.md for more details here.

Train MLC Model

Refer to toolrag/mlc/README.md for more details here.

Train ToolRefiner Model

Refer to toolrag/toolrefiner/README.md for more details here.

Citation

@misc{moon2024efficient,
      title={Efficient and Scalable Estimation of Tool Representations in Vector Space}, 
      author={Suhong Moon and Siddharth Jha and Lutfi Eren Erdogan and Sehoon Kim and Woosang Lim and Kurt Keutzer and Amir Gholami},
      year={2024},
      eprint={2409.02141},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2409.02141}, 
}

SqueezeAILab/Tool2Vec