Intro to AI Engineering with Python (4/6) – Rapid Prototyping and Demos Using Gradio

Hi and welcome back to part 4 of this tutorial series. In this part, we’ll look at using Gradio to quickly build an interface so we can have a prototype of whatever it is we’re building using LLMs, APIs, or even Machine Learning models.

This part will be slightly different from the others in that we’ll take a bit of time out to build a small project using everything we’ve learned so far and adding Gradio as a new component. The project we’ll build is a simple silly story generator and it should look like this at the end of this tutorial part:

We’ll have some fun with this so the project is admittedly a bit silly in nature, but the knowledge can be equally applied to serious matters and will show you the basics of quickly prototyping an idea or setup so you can share it with others who can then interact with it through a simple interface.

Getting started

The first step is to install Gradio, as it is simply a Python library. So run the following command in your terminal:

pip install gradio

Now that we have gradio installed, let’s get to it! Start by creating a new file in your project folder named gradio_project.py:

📁 Intro_to_AI_engineering
    📁 output
    📄 .env
    📄 chat_gpt_request.py
    📄 generate_image.py
    📄 generate_speech.py
    📄 gradio_project.py     ✨ New file
    📄 langchain_basics.py
    📄 test_audio.mp3
    📄 text_to_summarize.txt
    📄 transcribe_audio.py

And inside gradio_project.py, we’ll get started with our imports:

import gradio as gr
from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

These are all familiar imports from previous parts, except the gradio import. It’s a bit of a convention to use gr as the alias for gradio when importing it, so we’ll stick with that. Now that is nice and all, but as you can see from the project image I shared above, our prototype will also have an image.

Complete your imports by adding the following:

from generate_image import generate_image, download_and_save_image

load_dotenv()

We could also use LangChain to call Dall-e and generate our images that way, but in part 2 we already wrote reusable functions to generate an image and download it from the internet. We can simply import and reuse what we’ve already built before!

Finally, we call load_dotenv() as always, for the API key.

Writing the prompts

As we can see from the project image featuring Muscle Kiwi Man above, we are going to need a prompt for ChatGPT to write us a silly story on a particular topic. We’ll also need an image though and if you remember our generate_image function, it takes a description of the desired image to be generated as input.

If we have to write this description ourselves that would be very inconvenient, yet giving the whole short story to Dall-e will be too much conflicting information about what needs to be in the image. This is why we’ll need a second prompt for ChatGPT in which we’ll simply ask it to generate an appropriate description for the image for us, which can then go into our generate_image function.

Let’s start with the prompt for the story generation:

short_story_prompt = ChatPromptTemplate.from_template(
    """
    Please give me a short story where the main character is a(n) {word} and make it very bizarre. The story should have two to three paragraphs and proper formatting."
    Story:
    """
)

We use the same LangChain ChatPromptTemplate class we did before, making sure to put in a {word} placeholder for the main character of the story, which the user will provide for us. We also specify the desired length of the story and for it to have proper formatting (using paragraphs).

The ending just says Story: to invite the LLM to start writing the story for us right away, without any awkward introductions. Sometimes an LLM will answer like “Sure I can help you with that” and then start the story, but we don’t want that in our case. By stating Story: it wouldn’t make sense for ChatGPT to start with anything other than the story itself.

Next, we’ll write the prompt to ask for a good image description to fit the story:

image_description_prompt = ChatPromptTemplate.from_template(
    """
    Please take the following story and describe the image that would go with it (Return a prompt for an image generator). Keep in mind the main character of this story is a(n) {word}.
    Story: {story}
    Your prompt for the image generator:
    """
)

So here we just ask ChatGPT to describe an image that would go with the story we provide it, stating for clarity that we need a prompt for an image generator. We also remind it to keep the main character in mind as it should be important in the image. There are two placeholders for {word} and {story}, and we end with the Your prompt for the image generator: trick again to try and avoid any unnecessary introductions in the response.

Setting up our model and the Chains

Now we can set up our ChatGPT model and the Chains we’ll use to generate the story and the image description. Let’s define the llm and output_parser first:

llm = ChatOpenAI(model="gpt-4o-mini")
output_parser = StrOutputParser()

No surprises there, these are the same as we’ve used before. Again, feel free to use the larger gpt-4o model if you prefer. Now we’ll set up the individual chains for the story and the image description:

story_chain = short_story_prompt | llm | output_parser
image_description_chain = image_description_prompt | llm | output_parser

We simply used the | pipe syntax to chain the prompt into the LLM and then into the output parser for both chains. Again, if you want to learn more intricate and powerful ways to combine and create systems, refer to the LangChain/LangGraph course we have at the Finxter Academy. As this system is pretty simple, we’ll just use a function to call the chains in order:

def run_chain(word: str):
    story = story_chain.invoke({"word": word})
    print("story:", story, end="\n\n")

    image_description = image_description_chain.invoke(
        {
            "word": word,
            "story": story,
        }
    )
    print("image description:", image_description)

    image_url = generate_image(image_description)
    image_file_path = download_and_save_image(image_url)

    return story, image_file_path

So the entire function takes a single input of word which is a string value. This is the only thing we need from the user to run this. Technically, they can enter more than a single word (e.g. “Crazy Pineapple”), but we’ll just name it word for now. We then get the story by calling .invoke on the story_chain, passing in a dictionary with the key and value for word.

There is a print statement in there, just so we can also see what is going on in the terminal, and we then invoke the second chain as well, passing in both the word parameter and the story parameter which resulted from the last step. Another print statement purely for the terminal follows.

Now we call our generate_image function from part 2. But as Gradio will need a real image to display, and not just an internet URL, we also use our download_and_save_image function which will save the image for us. The function we wrote also returns the file path of the saved image which we save in a variable.

Finally, we have a double return of both the story and the image_file_path, which is all that Gradio will need to display the story and the image in our demo prototype.

Building the Gradio interface

Now that we have everything we want to run in our demo, it’s time to see how we can use Gradio to build a quick and simple user interface so that we can share our product prototype with less technical people.