Building Your First HuggingFace Transformers Tool

banjtheman

Banjo Obayomi

Posted on May 16, 2023

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
Enter fullscreen mode Exit fullscreen mode

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)
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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))
Enter fullscreen mode Exit fullscreen mode

Now, we can simply test the tool:

tool = CatImageFetcher()
tool()
Enter fullscreen mode Exit fullscreen mode

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])
Enter fullscreen mode Exit fullscreen mode

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")
Enter fullscreen mode Exit fullscreen mode

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")
Enter fullscreen mode Exit fullscreen mode

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

💖 💪 🙅 🚩
banjtheman
Banjo Obayomi

Posted on May 16, 2023

Join Our Newsletter. No Spam, Only the good stuff.

Sign up to receive the latest update from our blog.

Related