Update app.py

app.py

@@ -1,16 +1,17 @@
-# Import necessary libraries
 import streamlit as st
 from transformers import pipeline
 import torch
 from datasets import load_dataset
 
+# Ensure proper usage of device (CPU or GPU)
+device = 0 if torch.cuda.is_available() else -1  # Use GPU if available, else use CPU
+
 # Load the T5-based Emotion Classifier model
 @st.cache_resource
 def load_model():
     try:
         st.write("Loading the emotion analysis model...")
-        # Initialize the emotion classifier using the T5 model fine-tuned for emotion classification
-        emotion_analyzer = pipeline("text-classification", model="suryakiran786/T5-emotion")
+        emotion_analyzer = pipeline("text-classification", model="suryakiran786/T5-emotion", device=device)
         st.write("Model loaded successfully!")
         return emotion_analyzer
     except Exception as e:
@@ -20,24 +21,9 @@ def load_model():
 # Initialize the model (with caching to prevent reloads)
 emotion_analyzer = load_model()
 
-# Load the dataset if needed for any additional logic (not used in emotion analysis directly, just for example)
-@st.cache_data
-def load_data():
-    try:
-        # For demonstration purposes, let's load a sentiment analysis dataset from Hugging Face
-        dataset = load_dataset("glue", "sst2")
-        st.write("Dataset loaded successfully!")
-        return dataset
-    except Exception as e:
-        st.write(f"Error loading dataset: {e}")
-        return None
-
-# Load data (just to show usage, not used in emotion analysis directly)
-dataset = load_data()
-
 # Function to predict emotion for a single response
 def predict_emotion_single(response):
-    if not emotion_analyzer:
+    if emotion_analyzer is None:
         return {"Error": "Emotion analyzer model not initialized. Please check model loading."}
     try:
         response = response.strip()