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

# Emotions
emotions = ["Anger", "Love", "Fear", "Joy", "Sadness", "Surprise"]

# Load fine-tuned model
model_path = "./model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

def predict_emotions(comment):
    inputs = tokenizer(comment, return_tensors="pt", truncation=True)
    outputs = model(**inputs)
    scores = torch.sigmoid(outputs.logits)[0].detach().numpy()
    return {emotion: float(scores[i]) for i, emotion in enumerate(emotions)}

demo = gr.Interface(
    fn=predict_emotions,
    inputs=gr.Textbox(lines=4, placeholder="Enter GitHub comment here..."),
    outputs=gr.Label(num_top_classes=6),
    title="GitHub Comment Emotion Detector",
    description="Detects Anger, Love, Fear, Joy, Sadness, and Surprise in GitHub comments."
)

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