EmotionDetection-Text / src /streamlit_app.py
handex's picture
Update src/streamlit_app.py
2a7f672 verified
Raw
History Blame Contribute Delete
3.87 kB
import streamlit as st
import pickle
from collections import Counter
import re
from typing import Counter as TCounter, Tuple, List, Any
# --- 1. CONFIGURATION ---
# The n-gram range used during model training (e.g., (1, 4) for unigrams up to 4-grams)
N_GRAM_RANGE = (1, 4)
# Filenames saved in the previous step
MODEL_FILE = 'src/emotion.pkl'
VECTORIZER_FILE = 'src/vectorizer.pkl'
# --- 2. LOAD ARTIFACTS ---
@st.cache_resource
def load_artifacts():
"""Loads the saved classifier model and DictVectorizer."""
try:
# Load the Model
with open(MODEL_FILE, 'rb') as model_file:
loaded_model = pickle.load(model_file)
# Load the Vectorizer
with open(VECTORIZER_FILE, 'rb') as vec_file:
loaded_vectorizer = pickle.load(vec_file)
return loaded_model, loaded_vectorizer
except FileNotFoundError:
st.error(f"Error: Required files ({MODEL_FILE} or {VECTORIZER_FILE}) not found. Please ensure they are in the correct directory.")
return None, None
# --- 3. FEATURE EXTRACTION FUNCTION (CRITICAL!) ---
# This function MUST be identical to the one used during model training.
def create_feature(text: str, nrange: Tuple[int, int]) -> TCounter[str]:
"""Extracts n-gram and punctuation features from a text string."""
text_features: List[str] = []
lower_text = text.lower()
# Word N-Grams
text_alphanum = re.sub(r'[^a-z0-9#]', ' ', lower_text)
token = text_alphanum.split()
for n in range(nrange[0], nrange[1] + 1):
if n > 0 and n <= len(token):
ngrams = [' '.join(token[i-n:i]) for i in range(n, len(token) + 1)]
text_features.extend(ngrams)
# Punctuation Features
text_punc = re.sub(r'[a-z0-9]', ' ', lower_text)
text_features.extend(text_punc.split())
return Counter(text_features)
# --- 4. PREDICTION LOGIC ---
def predict_emotion(text: str, model: Any, vectorizer: Any) -> str:
"""Processes text and returns the predicted emotion label."""
if not text:
return "Please enter text for analysis."
# 1. Convert Text to Features (Counter object)
feature_counter = create_feature(text, N_GRAM_RANGE)
# 2. Vectorization (DictVectorizer expects a list of dictionaries/Counters)
# The saved vectorizer is used to ensure the features are ordered correctly.
X_processed = vectorizer.transform([feature_counter])
# 3. Make Prediction
prediction = model.predict(X_processed)
return prediction[0]
# --- 5. STREAMLIT UI ---
# Load the model artifacts
model, vectorizer = load_artifacts()
if model and vectorizer:
st.title("πŸ’¬ Text Emotion Detection App")
st.markdown("Enter a sentence or text below to see the predicted emotion.")
# User input area
user_input = st.text_area("Enter your text here:", "")
if st.button("Predict Emotion"):
if user_input:
# Perform prediction
result = predict_emotion(user_input, model, vectorizer)
# --- Result Display (Using visual cues) ---
st.subheader("Analysis Result")
# Simple color-coding based on common emotions
if result == 'joy':
st.success(f"Predicted Emotion: **{result.upper()}** πŸŽ‰")
elif result in ['sadness', 'fear']:
st.warning(f"Predicted Emotion: **{result.upper()}** 😟")
elif result == 'anger':
st.error(f"Predicted Emotion: **{result.upper()}** 😑")
elif result == 'disgust':
st.markdown(f"Predicted Emotion: **{result.upper()}** 🀒", unsafe_allow_html=True)
else:
st.info(f"Predicted Emotion: **{result.upper()}**")
else:
st.warning("Please enter some text to predict.")