Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import matplotlib.pyplot as plt | |
| import pandas as pd | |
| from sklearn.cluster import KMeans | |
| from sklearn.preprocessing import StandardScaler | |
| from transformers import pipeline | |
| import nltk | |
| import numpy as np | |
| from utils import read_poems_from_directory, emotion_labels_with_colors | |
| # Download nltk data for tokenization | |
| nltk.download('punkt') | |
| # Initialize emotion classifier pipelines | |
| models = { | |
| "Model 1": pipeline('sentiment-analysis', model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True), | |
| "Model 2": pipeline('sentiment-analysis', model="cardiffnlp/twitter-roberta-base-emotion", return_all_scores=True) | |
| } | |
| poems_directory = "./poems" | |
| os.makedirs(poems_directory, exist_ok=True) | |
| def analyze_poems_page(): | |
| st.header("Analyze Poems") | |
| # Sidebar for file upload and listing files | |
| st.sidebar.title("Upload New Poem") | |
| uploaded_file = st.sidebar.file_uploader("Choose a text file", type="txt") | |
| if uploaded_file is not None: | |
| with open(os.path.join(poems_directory, uploaded_file.name), "wb") as f: | |
| f.write(uploaded_file.getbuffer()) | |
| st.sidebar.success(f"Uploaded {uploaded_file.name}") | |
| st.sidebar.title("Available Poems") | |
| poem_files = [f for f in os.listdir(poems_directory) if f.endswith(".txt")] | |
| st.sidebar.write("\n".join(poem_files)) | |
| # Sidebar input for user-specified labels | |
| user_labels_input = st.sidebar.text_input("Enter emotion labels (comma-separated)", | |
| "happiness,sadness,anger,fear,disgust,surprise,love,joy,anxiety,contentment,frustration,loneliness,excitement,guilt,shame,envy,jealousy,pride,gratitude,empathy,compassion,boredom,relief,curiosity,awe,confusion,nostalgia,hope,despair,embarrassment") | |
| user_labels = [label.strip() for label in user_labels_input.split(",")] | |
| if st.button("Analyze Poems"): | |
| if os.path.isdir(poems_directory): | |
| # Read poems from the specified directory | |
| poems = read_poems_from_directory(poems_directory) | |
| if poems: | |
| def analyze_emotions(poem, model): | |
| lines = nltk.sent_tokenize(poem) | |
| emotions = [] | |
| for line in lines: | |
| result = model(line) | |
| emotions.append(result) | |
| return emotions | |
| def process_emotions(emotions): | |
| emotion_scores = [] | |
| all_labels = set() | |
| for line_emotions in emotions: | |
| line_score = {emo['label']: emo['score'] for emo in line_emotions[0]} | |
| all_labels.update(line_score.keys()) | |
| emotion_scores.append(line_score) | |
| return emotion_scores, all_labels | |
| def plot_emotional_arc(processed_emotions, labels, model_name): | |
| st.subheader(model_name) | |
| plt.figure(figsize=(15, 10)) | |
| for i, emotions in enumerate(processed_emotions): | |
| for emotion in labels: | |
| emotion_arc = [line_emotions.get(emotion, 0) for line_emotions in emotions] | |
| color = emotion_labels_with_colors.get(emotion, 'black') # default to black if not found | |
| plt.plot(emotion_arc, label=f'Poem {i+1} - {emotion}', color=color) | |
| plt.title(f'Emotional Arc of Each Poem ({model_name})') | |
| plt.xlabel('Line Number') | |
| plt.ylabel('Emotion Score') | |
| plt.legend(loc='center left', bbox_to_anchor=(1, 0.5)) | |
| st.pyplot(plt) | |
| def extract_features(emotion_data, labels): | |
| features = [] | |
| for emotions in emotion_data: | |
| poem_features = [] | |
| for label in labels: | |
| scores = [line_emotions.get(label, 0) for line_emotions in emotions] | |
| mean_score = np.mean(scores) | |
| std_score = np.std(scores) | |
| poem_features.extend([mean_score, std_score]) | |
| features.append(poem_features) | |
| return features | |
| # Analyze and plot for each model | |
| for model_name, model in models.items(): | |
| poem_emotions = [analyze_emotions(poem, model) for poem in poems] | |
| processed_emotions = [] | |
| all_labels = set() | |
| for emotions in poem_emotions: | |
| processed, labels = process_emotions(emotions) | |
| processed_emotions.append(processed) | |
| all_labels.update(labels) | |
| selected_labels = [label for label in user_labels if label in all_labels] | |
| plot_emotional_arc(processed_emotions, selected_labels, model_name) | |
| # Extract features for clustering | |
| features = extract_features(processed_emotions, selected_labels) | |
| # Create a DataFrame to store the features | |
| columns = [] | |
| for label in selected_labels: | |
| columns.extend([f'{label}_mean', f'{label}_std']) | |
| df = pd.DataFrame(features, columns=columns) | |
| # Standardize the features | |
| scaler = StandardScaler() | |
| scaled_features = scaler.fit_transform(df) | |
| # Apply KMeans clustering | |
| kmeans = KMeans(n_clusters=2, random_state=42) | |
| kmeans.fit(scaled_features) | |
| df['Cluster'] = kmeans.labels_ | |
| # Display the DataFrame | |
| st.write(f"Poem Sentiment Features and Clusters ({model_name}):") | |
| st.dataframe(df) | |
| # Visualize the clusters | |
| if not df.empty: | |
| plt.figure(figsize=(8, 6)) | |
| plt.scatter(df.iloc[:, 0], df.iloc[:, 1], c=df['Cluster'], cmap='viridis', marker='o') | |
| plt.title(f'Clusters of Poem Emotional Arcs ({model_name})') | |
| plt.xlabel(f'{columns[0]}') | |
| plt.ylabel(f'{columns[1]}') | |
| st.pyplot(plt) | |
| else: | |
| st.warning("No text files found in the specified directory.") | |
| else: | |
| st.error("The specified path is not a valid directory.") | |