Using & Mixing Hugging Face Models with Gradio 2.0

Using and Mixing Hugging Face Models with Gradio 2.0

The Hugging Face Model Hub has grown to become a vast repository of over 10,000 machine learning models submitted by users. This collection includes a wide range of natural language processing models that can perform tasks such as translating between Finnish and English or recognizing Chinese speech. More recently, the Hub has expanded to include models for image classification and audio processing.

Hugging Face has always aimed to make models accessible and easy to use. The transformers library allows users to load a model in just a few lines of code. Once a model is loaded, it can be used to make predictions on new data programmatically. However, machine learning models are not just used by programmers. An increasingly common scenario in machine learning is demoing models to interdisciplinary teams or letting non-programmers use models to help discover biases and failure points.

The Gradio library provides a simple way for machine learning developers to create demos and GUIs from machine learning models. These demos can be shared with collaborators as easily as sharing a Google Docs link. With the release of Gradio 2.0, users can now load and use almost any Hugging Face model with a GUI in just one line of code.

Example Usage

Here's an example of how to use a Hugging Face model with Gradio 2.0:

import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "distilbert-base-uncased-finetuned-sst-2-english"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    return outputs.logits

demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(label="Enter text to classify"),
    outputs=gr.Label(label="Classification"),
    title="Text Classification Demo",
    description="Classify text using the DistilBERT model",
)

demo.launch()

This code uses the AutoModelForSequenceClassification and AutoTokenizer classes from the transformers library to load the DistilBERT model and tokenizer. The classify_text function takes a text input, tokenizes it, and passes it through the model to get the classification output. The gr.Interface class is used to create a GUI demo that takes a text input and displays the classification output.

Customizing the Demo

Users can customize the demo by overriding any of the default parameters of the Interface class. For example, you can change the title and description of the demo, or add additional inputs or outputs.

Mixing Models

Gradio 2.0 also allows users to mix multiple models together to create more sophisticated applications. For example, you can load multiple models in parallel to compare their performance on a given task. Or, you can put models in series to create a pipeline that applies multiple models to a single input.

Here's an example of how to mix multiple models together:

import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model1_name = "distilbert-base-uncased-finetuned-sst-2-english"
model1 = AutoModelForSequenceClassification.from_pretrained(model1_name)
tokenizer = AutoTokenizer.from_pretrained(model1_name)

model2_name = "bert-base-uncased-finetuned-sst-2-english"
model2 = AutoModelForSequenceClassification.from_pretrained(model2_name)

def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs1 = model1(**inputs)
    outputs2 = model2(**inputs)
    return outputs1.logits, outputs2.logits

demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(label="Enter text to classify"),
    outputs=[gr.Label(label="Classification 1"), gr.Label(label="Classification 2")],
    title="Text Classification Demo",
    description="Classify text using two models",
)

demo.launch()