File size: 996 Bytes
4a2ada5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
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()
|