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

# โหลด tokenizer และโมเดล
emotion_model_name = "alexandrainst/da-emotion-classification-base"
emotion_tokenizer = AutoTokenizer.from_pretrained(emotion_model_name)
emotion_model = AutoModelForSequenceClassification.from_pretrained(emotion_model_name)

def main():
    st.title("Emotion Classification App")

    # สร้าง text area สำหรับป้อนข้อความ
    text = st.text_area('Enter the text for emotion classification:', '')

    # คลิกปุ่ม "Classify Emotion" เพื่อจำแนกอารมณ์ในข้อความ
    if st.button('Classify Emotion'):
        if text:
            # ใช้โมเดลเพื่อจำแนกอารมณ์ในข้อความ
            inputs = emotion_tokenizer(text, return_tensors="pt")
            outputs = emotion_model(**inputs)
            logits = outputs.logits

            # คำนวณคะแนนความน่าจะเป็นของแต่ละอารมณ์
            probabilities = logits.softmax(dim=1)

            # หาอารมณ์ที่มีความน่าจะเป็นสูงสุด
            predicted_emotion = torch.argmax(probabilities, dim=1).item()

            # แสดงอารมณ์ที่ทำนาย
            st.subheader("Predicted Emotion:")
            st.write(predicted_emotion)
        else:
            st.warning("Please enter some text for emotion classification.")

if __name__ == "__main__":
    main()