Merge remote-tracking branch 'origin/main'

app.py

@@ -1,33 +1,43 @@
-# Import necessary libraries
 import streamlit as st
-import transformers
 import torch
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Set up the Streamlit app
+# Set up Streamlit
 st.title("Emotion Detection with Transformers")
 
-# Create a text input widget
+# Text input
 user_input = st.text_area("Enter your text:")
 
 
-# Define a function for sentiment analysis using transformers
-@st.cache_data
-def load_model():
-    return pipeline("sentiment-analysis")
+# Function to load model and tokenizer using @st.cache_data
+@st.cache_data()
+def load_model_and_tokenizer():
+    model_name = "mrm8488/t5-base-finetuned-emotion"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return tokenizer, model
 
 
-# Load the sentiment analysis model
-sentiment_analyzer = load_model()
+tokenizer, model = load_model_and_tokenizer()
 
-# Create a button to analyze the emotion
+
+# Function to analyze emotion
+def analyze_emotion(text):
+    if text.strip() == "":
+        return "Please enter some text to analyze."
+
+    input_ids = tokenizer.encode(text + '</s>', return_tensors='pt')
+
+    output = model.generate(input_ids=input_ids,
+                            max_length=2)
+
+    dec = [tokenizer.decode(ids, skip_special_tokens=True) for ids in output]
+    label = dec[0]
+
+    return f"Emotion: {label.capitalize()}"
+
+
+# Analyze button
 if st.button("Analyze Emotion"):
-    if user_input:
-        # Perform sentiment analysis on user input
-        result = sentiment_analyzer(user_input)
-
-        # Display the result
-        emotion = result[0]['label']
-        st.write(f"Emotion: {emotion}")
-    else:
-        st.warning("Please enter some text to analyze.")
+    result = analyze_emotion(user_input)
+    st.write(result)