Hugging Face Course (3/6) – Creating an Interface for Our Chatbot

👉 Back to the Full Course on local models and Hugging Face (+Videos)

Welcome back to part 3! In this part, we will be using Gradio to create a quick but powerful interface for our chatbot. Gradio is a Python library that allows you to quickly create customizable UI components around your machine-learning models using the many pre-built components available.

To get started, let’s install Gradio in our virtual environment. Run the following command in a terminal window:

pipenv install gradio

Integrating our local chatbot with Gradio

Now it’s time to revisit our local_chat_memory.py file. We will be making some changes to it to integrate it with Gradio. Open local_chat_memory.py back up and get started by adding the following import to your imports up top:

import uuid
import threading
from typing import Iterator # <--- new

... (other imports) ...

We will use this as a type hint later on as we want to send an Iterator for Gradio to loop over.

Now let’s go to the constants below the imports, where I only changed the SYSTEM_MESSAGE constant:

store = {}
MODEL_NAME = "llama3"
LLAMA3 = Ollama(model=MODEL_NAME)
SYSTEM_MESSAGE = "You are a wise Jedi Master and helpful assistant. Your name is Yoda and you will address the user as Obi-Wan Kenobi at all times." # <--- updated
SESSION_ID = uuid.uuid4()

This is entirely optional but I felt like testing with a Star Wars theme this time because why not? 😄 Feel free to use a sensible system message instead.

Now scroll down to the chat function. We’ll be rewriting this function so that it has a normal mode which prints to the console and a Gradio mode that returns the iterator instead so that Gradio can loop over and display the messages. Here is the updated chat function:

def chat(query: str, gradio_mode: bool = False) -> None | Iterator[str]:
    response = llama3_w_memory.stream(
        {"query": query}, config={"configurable": {"session_id": SESSION_ID}}
    )
    if gradio_mode:
        return response
    else:
        for chunk in response:
            print(chunk, end="")

So we have our updated chat function that now takes an additional parameter gradio_mode which is a boolean that defaults to False. The function will either return None or an Iterator that will yield str values.

We call llama3_w_memory.stream like before with the query and config objects, catching the response. If Gradio mode is on we simply return this Iterator named response. If not, we loop over the response and print it to the console like the old function did, so we can still run this script in the terminal like we have done so far.

Now we’ll add a new function below the chat function where we put the model preload logic, taking it out of the if __name__ == "__main__": block. This is so we can reuse this function both in this script and in our Gradio interface script:

def preload(model_name: str = MODEL_NAME) -> None:
    thread = threading.Thread(target=preload_model, args=[model_name])
    thread.start()


if __name__ == "__main__":
    preload()
    try:
        while True:
            query = input("You: ")
            chat(query)
            print("\n")
    except KeyboardInterrupt:
        print("Shutting down...")

We have a new function called preload that takes a model_name parameter which defaults to our MODEL_NAME constant, so it doesn’t have to be set but you have the ability to overwrite it should you want to in the future. The function has no return. We then just call the thread logic that was in the if __name__ == "__main__": block before, and then we call the preload function from inside the if __name__ == "__main__": block instead. The rest of the block with the try and except is the same as before.

Feel free to test your local_chat_memory.py script in the terminal to make sure everything is still working as expected. It should behave exactly the same as before.

Adding the Gradio chat interface

Let’s move on to the Gradio chat app itself, which will be surprisingly simple to make. Gradio has a lower-level API using blocks that allow you to create custom interfaces for your models. This is useful if you want to create a more complex interface which is very customized to your model.

We will be using the higher-level API which has a pre-built interface for chatbots that we can use with minimal code, as this tutorial is not about Gradio. We don’t want to stray too far off-topic so we’ll just stick with quick-and-dirty!

Create a new file named chat_app.py in the root directory of your project:

📁Local_Models
    📁test_files
    📄chat_app.py   ✨New file
    📄local_chat.py
    📄local_chat_memory.py
    📄model_preloader.py
    📄Pipfile
    📄Pipfile.lock

Inside chat_app.py let’s start with our imports:

import gradio as gr
from local_chat_memory import chat, preload

No surprises here, we import Gradio and our chat and preload functions from local_chat_memory.py.

In order to use the prebuilt chat interface from Gradio we need to use the ChatInterface class. This ChatInterface class takes a function that we need to give as input that will be called whenever a message is sent.

The ChatInterface class will then send both the current query and the message history to this function that we give it each time the user makes a request.

ChatInterface (gradio) -> {query, history} -> our_function

Our function then in turn is supposed to return a response in the form of a generator, that can be looped over by Gradio to display the messages in its prebuilt interface.

our_function -> {response} -> ChatInterface (gradio)

Let’s first define our_function, and then afterward we can call the ChatInterface class and give it our function as input:

def run_gradio_chat(query, _):
    response = chat(query, gradio_mode=True)
    full_message = ""
    if response:
        for chunk in response:
            full_message += chunk
            yield full_message
    else:
        yield "There was an error fetching the response. Please try again."

We named our function run_gradio_chat and we take the query and history as input like shown above, except we will not be using the history parameter so we just named it _ for now. We already have the history we created in LangChain so we don’t need another history.

We then get the response by simply calling our chat function with the query and gradio_mode set to True so we get the Iterator back. If there is a response, we’re going to loop over each chunk in the response and append it to the full_message string, which starts as an empty string.

We then yield the full_message string which means that instead of waiting for the entire response to be assembled and then returned, the function run_gradio_chat can start providing parts of the response as soon as they are ready.

If there is no response, we yield a simple error message instead.

💡 In Python, yield is a keyword that is used in defining a generator. A generator is a special type of iterator. Instead of creating and returning a whole list of items at once (like a function would do), a generator uses the yield keyword to yield one item at a time, waiting to produce the next item until the current one has been consumed. This allows our chatbot here to ‘return’ what it has received so far bit by bit as it comes in.

Now let’s define an if __name__ == "__main__": block where we use our preload function and then create the Gradio interface:

if __name__ == "__main__":
    preload()
    gr.ChatInterface(run_gradio_chat).launch()

Preload will just start loading our llama 3 model in a separate thread without blocking the script just like we defined. Then we