Chat2VIS: AI-driven data visualisations with Streamlit, ChatGPT and natural language

Leverage large language models for Python code generation using prompt engineering

8 min readAug 1, 2023

Many exciting research projects are born out of our universities, and I’ve been privileged to be involved with some. But one thing that really drives me is bringing our research to life.

The release of ChatGPT in late 2022 inspired me and my co-researcher to explore how large language models (LLMs) could be utilised to generate data visualisations from natural language text. Nothing is more frustrating than hunting through menu items trying to find a command to change some plot aesthetic. Wouldn’t it be nice to use everyday language to graph what you would like to see?

So I decided to build Chat2VIS, to bring our research to you. You may have read my earlier post on this topic, but here I will delve a little further into the detail.

What is Chat2VIS?
How to use Chat2VIS
How to build Chat2VIS

Want to dive right in? Explore Chat2VIS, read our published research article, and take a look at the repo.

What is Chat2VIS?

Chat2VIS is an app that generates data visualisations via natural language using LLMs GPT-3 and ChatGPT (GPT-3.5 & GPT-4). There used to be Codex in there too earlier this year. You can ask it to visualise anything from the pre-loaded datasets, or load a dataset of your own.

Let me show how it works by using an example I’ve constructed.

Have you heard of speedcubing? In speedcubing competitions, competitors race to solve the Rubik’s Cube puzzle and beat their own personal best times. There are events for solving 3x3, 4x4, 5x5, 6x6, and 7x7 Rubik’s Cubes and many more —including blindfolded solving!

The competition results database is publicly available*, so I created a subset of it with results up to 23 June 2023. I took each competitor’s fastest best-solve time (as opposed to average-solve time) and I used the results from 2x2, 3x3, 4x4, 5x5, Clock, Megaminx, Pyraminx, Skewb, Square-1, and 3x3 blindfolded events. That’s 195,980 competitors in total — a dataset of 585,154 rows. Each row lists the competitor’s WCA ID, event name, best-solve time (in centiseconds), country, country ranking, continent ranking, and world ranking.

Here is what it looks like:

App overview

Let’s see how the app works:

Choose a pre-loaded dataset or upload one of your own.
Write the query in the language of your preference (no worries about spelling or grammar!)
Chat2VIS builds a unique prompt tailored to your dataset (the prompt template is generic enough so that each LLM understands the requirements without customisation).
Submit the prompt (which includes the beginnings of your Python script) to each LLM and get a continuation of your script (read more about it here).
Build the Python script by amalgamating the beginnings of the script from your initial prompt and the continuation script from the LLM.
Create the visualisation by rendering the script on the Streamlit interface. If you get no plot or a plot of something unexpected, it means the code has syntax errors (kind of like the code from human programmers!). Either give it another go and resubmit the query as is, or change your wording slightly before resubmitting the request.

How to use Chat2VIS

To begin, follow these steps:

Load the dataset.
Enter your OpenAI API key (if you don’t have one, get it here and add some credit).

Now you’re ready!

Example 1

Let’s start with a simple example from the speedcubing dataset.

Type in this query: “Show the number of competitors who have competed in the 3x3 event by country for the top 10 countries.”

Both ChatGPT-3.5 and GPT-3 performed well in understanding the query text and displaying the results, complete with axis labels and a title. They even correctly identified the “3x3 event” as the “3x3x3 Cube” value in the “Event Name” column. The USA has the highest number of speedcubers at approximately 38,000. However, ChatGPT could improve readability by changing the orientation of the x-axis bar labels. Maybe let the model know the preferred label orientation as a follow-up query.

Example 2

Let’s try a more challenging example.

Type in this query: “For each event, show the fastest best single time and put the value above the bar line. The results are in centiseconds. Please convert them to seconds.”

The LLMs are primarily trained in the English language but have knowledge of other languages as well.

Let’s add some multilingual text:

“Dessinez le tracé horizontalement” (“Draw the plot horizontal” in French)
“Whakamahia nga tae whero, kikorangi” (”Use red and blue colours” in te reo Māori, one of New Zealand’s official languages)

How did Chat2VIS do? Pretty good. The values are above the bar lines, the results are converted to seconds, the plot is turned horizontal, and the colours are red and blue. It even got the axis labels and the title right. Just look at that 3x3 time … 3.13 seconds! 👏

For more multilingual examples, queries with spelling mistakes, and refining of chart elements, read this article.

How to build Chat2VIS

Here is how to set up the front end:

To centre the titles and change the font, use st.markdown:

st.markdown("<h1 style='text-align: center; font-weight:bold; font-family:comic sans ms; padding-top: 0rem;'>Chat2VIS</h1>", unsafe_allow_html=True) 
st.markdown("<h2 style='text-align: center; padding-top: 0rem;'>Creating Visualisations using Natural Language with ChatGPT </h2>", unsafe_allow_html=True)

Create a sidebar and load the available datasets into a dictionary. Storing them in the session_state object avoids unnecessary reloading. Use radio buttons to select the chosen dataset, but also include any manually uploaded datasets in the list. To do this, add an empty container to reserve the spot on the sidebar, add a file uploader, and add the uploaded file to the dictionary. Finally, add the dataset list of radio buttons to the empty container (I like using emoji shortcodes on the labels!). If a dataset has been manually uploaded, ensure that the radio button is selected:

if "datasets" not in st.session_state:
    datasets = {}
    # Preload datasets
    datasets["Movies"] = pd.read_csv("movies.csv")
    datasets["Housing"] = pd.read_csv("housing.csv")
    datasets["Cars"] = pd.read_csv("cars.csv")
    st.session_state["datasets"] = datasets
else:
    # use the list already loaded
    datasets = st.session_state["datasets"]

with st.sidebar:
    # First we want to choose the dataset, but we will fill it with choices once we've loaded one
    dataset_container = st.empty()
    # Add facility to upload a dataset
    uploaded_file = st.file_uploader(":computer: Load a CSV file:", type="csv")
    # When we add the radio buttons we want to default the selection to the first
    index_no = 0
    if uploaded_file:
        # Read in the data, add it to the list of available datasets. Give it a nice name.
        file_name = uploaded_file.name[:-4].capitalize()
        datasets[file_name] = pd.read_csv(uploaded_file)
    # We want to default the radio button to the newly added dataset
    index_no = len(datasets)-1
    # Radio buttons for dataset choice
    chosen_dataset = dataset_container.radio(":bar_chart: Choose your data:", datasets.keys(), index=index_no)

Add checkboxes in the sidebar to choose which LLM to use. The label will display the model name with the OpenAI model version in brackets. The models and their selected status will be stored in a dictionary:

available_models = {"ChatGPT-4": "gpt-4", "ChatGPT-3.5": "gpt-3.5-turbo", "GPT-3": "text-davinci-003"}
with st.sidebar:
   st.write(":brain: Choose your model(s):")
   # Keep a dictionary of whether models are selected or not
   use_model = {}
   for model_desc,model_name in available_models.items():
      label = f"{model_desc} ({model_name})"
      key = f"key_{model_desc}"
      use_model[model_desc] = st.checkbox(label,value=True,key=key)

In the main section, add a password input widget for the OpenAI API key 🔑. The help parameter provides some extra information. Additionally, a text area for the query 👀 and a “Go” button are included.

my_key = st.text_input(label = ":key: OpenAI Key:", help="Please ensure you have an OpenAI API account with credit. ChatGPT Plus subscription does not include API access.", type="password")
question = st.text_area(":eyes: What would you like to visualise?", height=10)
go_btn = st.button("Go…")

Finally, display the datasets using a tab widget.

tab_list = st.tabs(df_list.keys()) 
for dataset_num, tab in enumerate(tab_list): 
    with tab: 
      dataset_name = list(df_list.keys())[dataset_num] 
      st.subheader(dataset_name) 
      st.dataframe(df_list[dataset_name], hide_index=True)

To initiate the process, click on “Go…”!

The communication with each model is facilitated through the openai Python library. With GPT-3, the prompt is presented as a sequence of tokens using the text completion endpoint API. ChatGPT models require the chat-completion endpoint and submission of a message sequence, which is then converted to tokens using ChatML (Chat Markup Language).

The following function illustrates this process, taking parameters for the prompt (question_to_ask), the model type (gpt-4, gpt-3.5-turbo, or text-davinci-003), and your OpenAI key. This function is placed within a try block with except statements to capture any errors returned from the LLMs (read more here):

def run_request(question_to_ask, model_type, key):
   openai.api_key = key
   if model_type == "gpt-4" or model_type == "gpt-3.5-turbo":
       # Run ChatGPT API
       response = openai.ChatCompletion.create(
         model=model_type,
         messages=[
         {"role":"system", "content":"Generate Python Code Script."},
         {"role":"user", "content":question_to_ask}])
       res = response["choices"][0]["message"]["content"]
   else:
       response = openai.Completion.create(
         engine=model_type,
         prompt=question_to_ask,
         temperature=0,
         max_tokens=500,
         top_p=1.0,
         frequency_penalty=0.0,
         presence_penalty=0.0,
         stop=["plt.show()"])
       res = response["choices"][0]["text"] 
 return res

Dynamically create as many columns on the interface as you have models selected:

model_list = [model_name for model_name, choose_model in use_model.items() if choose_model]
if len(model_list) > 0:
   plots = st.columns(len(model_list))

After executing the final scripts, the results for each model are passed to the st.pyplot chart elements for rendering in the columns on the interface.

Wrapping up

I’ve showed you how to create a natural language interface that displays data visualisations using everyday language requests on a set of data. I didn’t cover the details of engineering the prompt for the LLMs, but the referenced articles should give you more guidance and understanding. Since the development of Chat2VIS in January 2023, there have been significant advancements leveraging generative AI for visualisations and prompt engineering. There is so much more to explore!

Thank you to my research co-author and to those of you who have contacted me to show me how you have used it with your own datasets. It’s awesome to see! Thanks for reading and feel free to connect with me on LinkedIn. (Edited blog originally published on Streamlit )

*This information is based on competition results owned and maintained by the World Cube Association, published at https://worldcubeassociation.org/export/results as of June 23, 2023.