import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load the model and tokenizer
model_name = "WhoLetMeCook/ChefBERT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Function to make predictions
def predict_emotion(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=-1).item()
    return "Positive Emotion" if prediction == 1 else "Negative Emotion"

# Streamlit app layout
st.title("ChefBERT Emotion Classifier")
st.write("Enter a sentence and ChefBERT will predict whether the emotion is positive or negative.")

# Input box
user_input = st.text_input("Input Sentence", "")

if user_input:
    result = predict_emotion(user_input)
    st.write("Prediction: ", result)