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()