|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
raw_model_name = 'distilbert-base-uncased' |
|
raw_model = pipeline('sentiment-analysis', model=raw_model_name) |
|
fine_tuned_model_name = 'distilbert-base-uncased-finetuned-sst-2-english' |
|
fine_tuned_model = pipeline('sentiment-analysis', model=fine_tuned_model_name) |
|
|
|
def get_model_output(input_text, model_choice): |
|
raw_result = raw_model(input_text) |
|
fine_tuned_result = fine_tuned_model(input_text) |
|
return format_model_output(raw_result[0]), format_model_output(fine_tuned_result[0]) |
|
|
|
def format_model_output(output): |
|
return f"I am {output['score']*100:.2f}% sure that the sentiment is {output['label']}" |
|
|
|
iface = gr.Interface( |
|
fn=get_model_output, |
|
title="DistilBERT Sentiment Analysis", |
|
inputs=[ |
|
gr.Textbox(label="Input Text"), |
|
], |
|
outputs=[ |
|
gr.Textbox(label="Base DistilBERT output (distilbert-base-uncased)"), |
|
gr.Textbox(label="Fine-tuned DistilBERT output") |
|
], |
|
) |
|
|
|
iface.launch() |
|
|