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