import gradio as gr
from transformers import pipeline
import os
from dotenv import load_dotenv

load_dotenv()
token = os.environ.get("HF_TOKEN")

raw_model_name = "distilbert-base-uncased"
raw_model = pipeline("sentiment-analysis", model=raw_model_name)

scademy_500_500_model_name = "scademy/DistilBert-500-500-0"
scademy_500_500_model = pipeline("sentiment-analysis", model=scademy_500_500_model_name, token=token)

scademy_2500_2500_model_name = "scademy/DistilBERT-2500-2500-0"
scademy_2500_2500_model = pipeline("sentiment-analysis", model=scademy_2500_2500_model_name, token=token)

scademy_12500_12500_model_name = "scademy/DistilBERT-12500-12500-0"
scademy_12500_12500_model = pipeline("sentiment-analysis", model=scademy_12500_12500_model_name, token=token)

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)
    scademy_500_500_result = scademy_500_500_model(input_text)
    scademy_2500_2500_result = scademy_2500_2500_model(input_text)
    scademy_12500_12500_result = scademy_12500_12500_model(input_text)
    fine_tuned_result = fine_tuned_model(input_text)
    return (
        format_model_output(raw_result[0]),
        format_model_output(scademy_500_500_result[0]),
        format_model_output(scademy_2500_2500_result[0]),
        format_model_output(scademy_12500_12500_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="Scademy DistilBERT 500-500 output"),
        gr.Textbox(label="Scademy DistilBERT 2500-2500 output"),
        gr.Textbox(label="Scademy DistilBERT 12500-12500 output"),
        gr.Textbox(label="Fine-tuned DistilBERT output"),
    ],
)

iface.launch()