import gradio as gr
from fastai.text.all import load_learner

# Load FastAI model once
learn = load_learner("emotion_classifier.pkl")
learn.push_to_hub("fastai-emotion-classifier") 
learn = load_learner("https://huggingface.co/haripriyaram/fastai-emotion-classifier/resolve/main/export.pkl")

# Prediction function
def predict_emotion(text):
    pred_label, _, probs = learn.predict(text)
    probs_dict = {label: float(prob) for label, prob in zip(learn.dls.vocab, probs)}
    return pred_label, probs_dict

# Gradio UI
iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence..."),
    outputs=[
        gr.Label(label="Predicted Emotion"),
        gr.JSON(label="Confidence Scores")
    ],
    title="🎭 Emotion Classifier (FastAI)",
    description="Enter a sentence and the model will predict the corresponding emotion.",
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()