Gradio — Best I/O Framework for AI Models

Vito Hsu
Stackademic
Published in
3 min readMay 7, 2024

--

Before I started using the Grafio framework, I used the Streamlit framework.

However, when I first used the Grafio framework a little more deeply, I found that it might be lighter than the Streamlit framework. What I mean is that it seems that there is no need to care about the Component. For Machine Learning Engineers, this can reduce the design thinking of the UI as much as possible and focus on the accuracy of the Model.

In addition to the need to improve the accuracy of the model, as an engineer, more importantly, we often face the boss or supervisor on the front line, so we need a simple and fast Demo tool. So that they can quickly understand the current situation of the model.

Different from the common simple examples of Gradio. In this article, I want to explain how to quickly implement the input/output design of multiple prediction models on Gradio. The key point is that multiple models are used in the same application with little code.

All of the following models are ready-made from the Transformers package, so we don’t need to train any model, we just need to focus on which type of model to use. This is also very important for us, because in the future, perhaps in the field of AI, what we need to understand more is not the design details of all models, but how to use the right models at the right time.

在我開始使用 Grafio 框架之前,我原先常用的是 Streamlit 框架。

然而,當我第一次稍微深度使用了解 Gradio 框架後,我發現它可能比起 Streamlit 框架來說更輕量,我的意思是,似乎不太需要去管 Component。這對於 Machine Learning Engineer 來說,能減少 UI 的設計思維,就盡量降低,專注在 Model 的精確度。

除了模型的準確度需要提升之外,身為工程師,更重要的是,常常第一線面對的是老闆或主管,所以,我們更需要一個簡單快速的 Demo 工具。好讓它們快速理解這個模型目前的情況。

有別於 Gradio 常見的簡單範例。本篇,我想要說明的是,如何在 Gradio 上面快速實現多個預測模型之輸入/輸入的設計。重點在於多個模型用在同一個應用程式,且使用的代碼極少。

以下所有模型都從 Transformers 套件裡面拿現成的,所以我們不需要去訓練,我們只需要專注於要使用哪類型的模型。這件事情對我們來說,也是很關鍵的,原因是未來或許 AI 的領域裡面,我們更需要了解的並非所有模型的設計細節,而是如何在對的時候用上對的工具。

Code

For Only One Model


import gradio as gr
from transformers import pipeline

# Load Model
classifier = pipeline("text-classification", model="distilbert-base-uncased")

def classify_text(text):
result = classifier(text)
label = result[0]['label']
score = result[0]['score']
return f"Label: {label}, Confidence (0-1): {score:.4f}"

iface = gr.Interface(
fn=classify_text,
inputs=gr.Textbox(lines=5, max_lines=100, label="Enter Your Text here !!"),
outputs="text",
title="text classification",
description="positive / negative",
theme="compact"
)

iface.launch()

For Multiple Model (Here we demo three models together.)

import gradio as gr
from transformers import pipeline

classifier = pipeline("text-classification", model="distilbert-base-uncased")

text_generator = pipeline("text-generation", model="gpt2")

summarizer = pipeline("summarization") # default model="t5-base"

def classify_text(text):
result = classifier(text)
label = result[0]['label']
score = result[0]['score']
return f"Label: {label}, Confidence: {score:.4f}"

def generate_text(prompt):
generated_text = text_generator(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']
return generated_text

def summarize_text(text):
summarized_text = summarizer(text)[0]['summary_text']
return summarized_text

def combind(text1, prompt, text2):
return classify_text(text1), generate_text(prompt), summarize_text(text2)

with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
input1 = gr.Text(label="Input Area for Classification")
output1 = gr.Text(label="Output Area for Classification")
with gr.Column():
input2 = gr.Text(label="Input Area for Generation")
output2 = gr.Text(label="Output Area for Generation")
with gr.Column():
input3 = gr.Text(label="Input Area for Summarization")
output3 = gr.Text(label="Output Area for Summarization")
btn = gr.Button("Start")
btn.click(
combind,
inputs=[input1, input2, input3], # key1 : inputs
outputs=[output1, output2, output3] # key2 : outputs
)

demo.launch()
Result

SEE YA~

Stackademic 🎓

Thank you for reading until the end. Before you go:

--

--