import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

MODEL_ID = "tharakap/deberta-sentiment-analyser" 

CUSTOM_ID_TO_LABEL = {
    0: "anger",
    1: "fear",
    2: "joy",
    3: "sadness",
    4: "surprise"
}

def load_model_and_pipeline():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_ID, 
        id2label=CUSTOM_ID_TO_LABEL
    )
    
    return pipeline("text-classification", model=model, tokenizer=tokenizer, return_all_scores = True) 
try:
    classifier = load_model_and_pipeline()
except Exception as e:
    print(f"Error loading model: {e}")
    classifier = pipeline("sentiment-analysis") 


def predict_sentiment(text, threshold=0.2): # You set the threshold here
    results = classifier(text)[0]
    
    formatted = []
    active_labels = []

    for res in results:
        label = res["label"].upper()
        score = res["score"]
        
        # Check if the score is above the custom threshold
        if score >= threshold:
            active_labels.append(label)
            formatted.append(f"- **{label}** (ACTIVE): {round(score, 4)}")
        else:
            formatted.append(f"- {label}: {round(score, 4)}")
            
    output = "### Active Emotions (Threshold > " + str(threshold) + "):\n"
    if active_labels:
        output += ", ".join(active_labels) + "\n\n"
    else:
        output += "None detected above threshold.\n\n"
        
    output += "### All Emotion Scores:\n" + "\n".join(formatted)
    
    return output



iface = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(
        lines=5, 
        placeholder="Enter text for emotion analysis...",
        label="Input Text"
    ),
    outputs="markdown",
    title="DeBERTa Emotion Analysis Demo",
    description=f"An emotion classification demo using the DeBERTa model: {MODEL_ID}.",
    examples=[
        ["I am thrilled to hear the good news!"],
        ["I am terrified; the sight of the spider filled me with pure dread."],
        ["He is absolutely furious and enraged right now."] 
    ]
)

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