Building Your First HuggingFace Transformers Tool
Banjo Obayomi
Posted on May 16, 2023
In this post, we'll walk you through the process of building a simple tool using the Transformers library.
The tool we're building will fetch an image of a cat from the internet each time it is run. We'll be using the Cataas (Cat-As-A-Service) API to achieve this. Let's dive in!
Prerequisites
Before we embark on this exciting journey of building the Cat Fetcher tool, we need to ensure that our development environment is properly set up. This tool relies on several Python packages - requests
for making HTTP requests, Pillow
for handling images, and transformers
for creating the AI tool.
To install these packages, you can use the following command in your terminal:
pip install requests Pillow transformers
Fetching a Random Cat Image
To get started, here is the code that fetches a random cat image from the Cataas API:
import requests
from PIL import Image
image = Image.open(requests.get('https://cataas.com/cat', stream=True).raw)
Creating a Tool
To create a tool that can be used by our system, we first create a class that inherits from the superclass Tool:
from transformers import Tool
class CatImageFetcher(Tool):
pass
This class needs the following attributes:
name: This is the name of the tool. For consistency, we'll name it cat_fetcher.
description: This will be used to populate the prompt of the agent.
inputs and outputs: These help the interpreter make educated choices about types and allow for a demo to be spawned when we push our tool to the Hub. They are both a list of expected values, which can be text, image, or audio.
A call method: This contains the inference code.
So, our class now looks like this:
from transformers import Tool
from huggingface_hub import list_models
class CatImageFetcher(Tool):
name = "cat_fetcher"
description = ("This is a tool that fetches an actual image of a cat online. It takes no input, and returns the image of a cat.")
inputs = []
outputs = ["image"]
def __call__(self):
return Image.open(requests.get('https://cataas.com/cat', stream=True).raw).resize((256, 256))
Now, we can simply test the tool:
tool = CatImageFetcher()
tool()
Using the Tool with an Agent
To use the tool with an agent, we recommend instantiating the agent with the tools directly:
from transformers.tools import HfAgent
agent = HfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder", additional_tools=[tool])
Then, we can have the agent use it with other tools:
agent.run("Fetch an image of a cat online and caption it for me")
The tool successfully fetches a cat image, and the image captioning tool captions the image. That's it! You've just created your first tool with the Transformers library.
As a final step, we recommend pushing the tool to the Hugging Face Model Hub so others can benefit from it.
tool.push_to_hub("YOUR-TOOL")
Thank you for following along! We're excited to see the creative and useful tools you'll build.
To learn more about building tools check out the official documentation
Posted on May 16, 2023
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.