import gradio as gr
from huggingface_hub import from_pretrained_fastai

# Load model directly from Hugging Face
learn = from_pretrained_fastai("haripriyaram/Text-emotion-Recognizer-Model")

# Prediction function
def predict_emotion(text):
    pred_label, _, probs = learn.predict(text)

    # Handle nested vocab if needed
    vocab = learn.dls.vocab[0] if isinstance(learn.dls.vocab[0], list) else learn.dls.vocab
    probs_dict = {label: float(prob) for label, prob in zip(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 (ULMFit ML Service Deployment)",
    description="Enter a sentence and the model will predict the corresponding emotion.",
    allow_flagging="never"
)

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