Update app.py

app.py

@@ -1,153 +1,86 @@
-import streamlit as st
-import torch
-import pandas as pd
-import matplotlib.pyplot as plt
-from BertEmotionClassifier import BertEmotionClassifier
-from transformers import AutoTokenizer
-from huggingface_hub import hf_hub_download
-import torch.nn.functional as F
-
-# ----------------------------
-# Config
-# ----------------------------
-EMOTIONS = ['anger', 'fear', 'joy', 'sadness', 'surprise']
-MODEL_NAME = "roberta-base"
-
-st.set_page_config(page_title="Emotion Classifier", page_icon="🎭", layout="wide")
-
-# ----------------------------
-# Custom Styling
-# ----------------------------
-st.markdown("""
-    <style>
-        .emotion-card {
-            padding: 15px;
-            border-radius: 18px;
-            text-align: center;
-            font-size: 18px;
-            font-weight: 600;
-            color: white;
-            margin-bottom: 10px;
-            box-shadow: 0px 0px 10px rgba(0,0,0,0.4);
-        }
-        .dominant {
-            background: linear-gradient(135deg, #00c6ff, #0072ff);
-        }
-        .sub {
-            background: linear-gradient(135deg, #434343, #000000);
-        }
-        .main-container {
-            max-width: 900px;
-            margin: auto;
-            padding-top: 30px;
-        }
-    </style>
-""", unsafe_allow_html=True)
-
-# ----------------------------
-# Load Model
-# ----------------------------
-@st.cache_resource
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-
-    model = BertEmotionClassifier(model_name=MODEL_NAME, num_labels=len(EMOTIONS))
-
-    # Download model weights from HuggingFace Hub
-    model_path = hf_hub_download(repo_id="aadhi3/RoBert_Model", filename="model.pth")
-
-    # Load state dict from downloaded path
-    state_dict = torch.load(model_path, map_location="cpu")
-    state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
-    model.load_state_dict(state_dict)
-    model.eval()
-
-    return tokenizer, model
-
-tokenizer, model = load_model()
-
-# ----------------------------
-# Prediction Function
-# ----------------------------
-def predict_emotions(text: str):
-    encoding = tokenizer(
-        text,
-        add_special_tokens=True,
-        max_length=256,
-        padding='max_length',
-        truncation=True,
-        return_tensors='pt'
-    )
-
-    with torch.no_grad():
-        logits = model(**encoding)
-        probs = F.softmax(logits, dim=-1)[0].cpu()
-
-    return {emo: round(float(probs[i]), 4) for i, emo in enumerate(EMOTIONS)}
-
-# ----------------------------
-# UI Layout
-# ----------------------------
-st.title("🎭 Emotion Detection AI")
-st.markdown("### Understand emotions in text using **AI-driven emotion analysis**")
-st.write("")  # spacing
-
-input_text = st.text_area(
-    "Enter your text here:",
-    placeholder="Type something like: 'I am extremely happy today!' 😄",
-    height=150
-)
-
-predict_btn = st.button("🔮 Analyze Emotions", use_container_width=True)
-
-if predict_btn and input_text.strip() != "":
-    with st.spinner("AI thinking..."):
-        results = predict_emotions(input_text.strip())
-
-        df = pd.DataFrame(results.items(), columns=["Emotion", "Probability"])
-        dominant = df.iloc[df["Probability"].idxmax()]["Emotion"]
-
-        st.markdown(f"### 🎯 Dominant Emotion: **{dominant.upper()}**")
-
-        col_chart, col_cards = st.columns([1.2, 1])
-
-        # ---------------------------
-        # Matplotlib Bar Chart
-        # ---------------------------
-        with col_chart:
-            fig, ax = plt.subplots()
-            ax.bar(df["Emotion"], df["Probability"])
-            ax.set_ylim(0, 1)
-            ax.set_ylabel("Probability Score")
-            ax.set_title("Emotion Prediction Distribution")
-            st.pyplot(fig)
-
-        # ---------------------------
-        # Emotion Cards
-        # ---------------------------
-        with col_cards:
-            st.markdown("### 📊 Emotion Strength")
-            for emo, val in results.items():
-                style = "dominant" if emo == dominant else "sub"
-                st.markdown(
-                    f"<div class='emotion-card {style}'>{emo.capitalize()} — {val}</div>",
-                    unsafe_allow_html=True
-                )
-
-        # Save Prediction History
-        row = {"text": input_text, **results}
-        try:
-            old_df = pd.read_csv("./predictions.csv")
-        except FileNotFoundError:
-            old_df = pd.DataFrame(columns=["text"] + EMOTIONS)
-
-        new_df = pd.concat([old_df, pd.DataFrame([row])], ignore_index=True)
-        new_df.to_csv("./predictions.csv", index=False)
-
-        with st.expander("📂 View Recent Predictions"):
-            st.dataframe(new_df.tail(10), use_container_width=True)
-
-        st.success("Result saved successfully! ✨")
-
-st.markdown("---")
-st.caption("Built with ❤️ using Streamlit & PyTorch — deployed on HuggingFace Spaces")
+import streamlit as st
+import torch
+import pandas as pd
+import matplotlib.pyplot as plt
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
+from BertEmotionClassifier import BertEmotionClassifier
+
+# ----------------------------
+# Config
+# ----------------------------
+EMOTIONS = ['anger', 'fear', 'joy', 'sadness', 'surprise']
+MODEL_NAME = "roberta-base"
+
+st.set_page_config(page_title="Emotion Classifier", page_icon="🎭", layout="centered")
+
+st.title("🎭 Emotion Detection AI")
+st.markdown("Predict emotional sentiment from text using a fine-tuned RoBERTa model.")
+
+# ----------------------------
+# Load Model
+# ----------------------------
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = BertEmotionClassifier(model_name=MODEL_NAME, num_labels=len(EMOTIONS))
+
+    model_path = hf_hub_download(repo_id="aadhi3/RoBert_Model", filename="model.pth")
+
+    state_dict = torch.load(model_path, map_location="cpu")
+    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+    model.load_state_dict(state_dict)
+    model.eval()
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# ----------------------------
+# Prediction Function
+# ----------------------------
+def predict_emotions(text: str):
+    encoding = tokenizer(
+        text,
+        add_special_tokens=True,
+        max_length=256,
+        padding="max_length",
+        truncation=True,
+        return_tensors="pt"
+    )
+    with torch.no_grad():
+        logits = model(**encoding)
+        probs = F.softmax(logits, dim=-1)[0].cpu()
+
+    return {emo: round(float(probs[i]), 4) for i, emo in enumerate(EMOTIONS)}
+
+# ----------------------------
+# UI Layout
+# ----------------------------
+input_text = st.text_area("Enter text to analyze:", height=120)
+
+if st.button("🔮 Analyze"):
+    if input_text.strip():
+        with st.spinner("Analyzing emotions..."):
+            results = predict_emotions(input_text.strip())
+            df = pd.DataFrame(results.items(), columns=["Emotion", "Probability"])
+            dominant = df.loc[df["Probability"].idxmax(), "Emotion"]
+
+        st.markdown(f"### 🎯 Dominant Emotion: **{dominant.upper()}**")
+
+        # Chart
+        fig, ax = plt.subplots(figsize=(5, 3))
+        ax.bar(df["Emotion"], df["Probability"])
+        ax.set_ylim(0, 1)
+        ax.set_ylabel("Probability")
+        ax.set_title("Emotion Prediction")
+        st.pyplot(fig)
+
+        st.markdown("### 📊 Prediction Details")
+        cols = st.columns(len(EMOTIONS))
+
+        for col, (emo, val) in zip(cols, results.items()):
+            col.metric(label=emo.capitalize(), value=f"{val}")
+
+st.markdown("---")
+st.caption("Built with ❤️ using Streamlit & PyTorch — deployed on Hugging Face Spaces")