Spaces:
Sleeping
Sleeping
| 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 --- | |
| 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.") |