Building Your First HuggingFace Transformers Tool

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:

1. name: This is the name of the tool. For consistency, we'll name it cat_fetcher.

2. description: This will be used to populate the prompt of the agent.

3. 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.

4. 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

Comments

[adriens]

How would you deal with a API that requires header authentication tokens ?

[Banjo Obayomi]

requests.kennethreitz.org/en/lates...

[adriens]

Yep, I was aware of that. What I meant is : where is the optimal choice to configure the Bearer tokens ? ... Probably within env variables 💭 to match OPENAI practices ;-p

[Banjo Obayomi]

Yea, that sounds good!

[adriens]

Thanks a lot for having shared this very cool tutorial... and your use case around API cats is just very fun and very useful !