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


pipe = pipeline("text-classification", model="s1143700/Internship")

tokenizer = AutoTokenizer.from_pretrained("s1143700/Internship")
model = AutoModelForSequenceClassification.from_pretrained("s1143700/Internship")


def preprocess_text(text):
    inputs = tokenizer(
        text,
        truncation=True,
        padding="max_length",
        max_length=100,
        return_tensors="pt"
    )
    return inputs

def predict_semantics(text):
    # Tokenize
    inputs = preprocess_text(text)
    
    # Forward pass
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get probabilities
    logits = outputs.logits
    probabilities = torch.softmax(logits, dim=1)
    
    # Create a label dictionary for Gradio
    emotion_labels = ["sadness", "joy", "love", "anger", "fear", "surprise"]
    result = {
        emotion_labels[i]: float(probabilities[0][i])
        for i in range(len(emotion_labels))
    }
    
    return result

# Create the interface
iface = gr.Interface(
    fn=predict_semantics,
    inputs=gr.Textbox(label="Input Text", placeholder="Enter your text here..."),
    outputs=gr.Label(label="Emotion Probabilities"),
    title="Emotion Analysis",
    description="Enter text to analyze its emotional content (sadness, joy, love, anger, fear, surprise).",
    examples=[
        ["I'm so happy today!"],
        ["This situation makes me anxious"],
        ["I feel loved by my family"],
        ["That movie scared me"]
    ]
)

# Launch the interface
iface.launch()