Update app.py

app.py

@@ -62,19 +62,24 @@ RHETORICAL_DEVICES = {
 
 class SpeechAnalyzer:
     def __init__(self):
-    # Load MoralFoundations model
+        # Load MoralFoundations model
         self.moral_model_path = "MMADS/MoralFoundationsClassifier"
         self.moral_tokenizer = RobertaTokenizer.from_pretrained(self.moral_model_path)
         self.moral_model = RobertaForSequenceClassification.from_pretrained(self.moral_model_path)
     
-    # Define label names directly
+        # Define label names directly
         self.label_names = ['care', 'fairness', 'loyalty', 'authority', 'sanctity']
-        
+            
         # Other pipelines remain the same
         self.sentiment_pipeline = pipeline("sentiment-analysis")
         self.ner_tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
         self.ner_model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
         self.ner_pipeline = pipeline("ner", model=self.ner_model, tokenizer=self.ner_tokenizer)
+        
+        # Add emotion classifier
+        self.emotion_classifier = pipeline("text-classification", 
+                                         model="j-hartmann/emotion-english-distilroberta-base")
+
 
     def split_text(self, text, max_length=256, overlap=50):
         """Split long text into overlapping segments"""
@@ -335,68 +340,187 @@ def main():
                 st.write(f"**{MORAL_FOUNDATIONS[foundation]}**: {score:.2%}")
         
         with tab2:
+            status_text.text('Processing Emotional Trajectory...')
+            progress_bar.progress(40)
             st.subheader("Speech Trajectory Analysis")
-            col1, col2, col3 = st.columns(3)
             
+            compare_mode = st.toggle("Enable Side-by-Side Comparison", value=False)
+            
+            viz_options = st.multiselect(
+                "Select visualizations to display:",
+                ["Sentiment Flow", "Moral Foundations Flow", "Basic Emotions Flow"],
+                default=["Sentiment Flow"]
+            )
+            
+            if len(viz_options) > 1 and compare_mode:
+                cols = st.columns(len(viz_options))
+            else:
+                cols = [st]  # Single column mode
+            
+            # Create consistent segments for analyses
             segments = analyzer.split_text(text, max_length=512)
             num_segments = len(segments)
             segment_labels = [f"{i+1}" for i in range(num_segments)]
             
-            sentiment_scores, basic_emotions = analyzer.analyze_emotional_trajectory(text)
+            current_col = 0
             
-            with col1:
-                st.write("### Sentiment Flow")
-                sentiment_fig = go.Figure(data=go.Scatter(
-                    x=segment_labels,
-                    y=sentiment_scores,
-                    mode='lines+markers',
-                    line=dict(color='#1f77b4', width=3),
-                    marker=dict(
-                        size=8,
-                        color=['#ff0000' if score < 0.4 else '#00ff00' if score > 0.6 else '#888888' for score in sentiment_scores],
-                        symbol='circle'
+            if "Sentiment Flow" in viz_options:
+                with cols[current_col]:
+                    sentiment_scores = analyzer.analyze_emotional_trajectory(text)
+                    
+                    trajectory_fig = go.Figure(data=go.Scatter(
+                        x=segment_labels,
+                        y=sentiment_scores,
+                        mode='lines+markers',
+                        line=dict(color='#1f77b4', width=3),
+                        marker=dict(
+                            size=8, 
+                            color=['#ff4444' if score < -0.3 else '#44ff44' if score > 0.3 else '#888888' for score in sentiment_scores],
+                            symbol='circle'
+                        )
+                    ))
+                    
+                    trajectory_fig.update_layout(
+                        title='Emotional Flow Throughout the Speech',
+                        xaxis_title='Speech Segments',
+                        yaxis_title='Emotional Tone',
+                        yaxis=dict(
+                            ticktext=['Very Negative', 'Negative', 'Neutral', 'Positive', 'Very Positive'],
+                            tickvals=[-1, -0.5, 0, 0.5, 1],
+                            range=[-1.1, 1.1],
+                            gridcolor='lightgray'
+                        ),
+                        hovermode='x unified',
+                        showlegend=False,
+                        plot_bgcolor='white',
+                        height=500
                     )
-                ))
-                
-                sentiment_fig.update_layout(
-                    title='Sentiment Throughout Speech',
-                    xaxis_title='Speech Segments',
-                    yaxis_title='Sentiment',
-                    yaxis=dict(
-                        ticktext=['Very Negative', 'Negative', 'Neutral', 'Positive', 'Very Positive'],
-                        tickvals=[0, 0.25, 0.5, 0.75, 1],
-                        range=[0, 1],
-                        gridcolor='lightgray'
-                    ),
-                    plot_bgcolor='white'
-                )
-                st.plotly_chart(sentiment_fig)
-            
-            # Rest of the code remains the same for col2 (Moral Foundations)
+                    
+                    trajectory_fig.update_traces(
+                        hovertemplate="Segment: %{x}<br>Score: %{y:.2f}<extra></extra>"
+                    )
+                    
+                    st.plotly_chart(trajectory_fig, use_container_width=True)
+                    if compare_mode:
+                        current_col += 1
             
-            with col3:
-                st.write("### Basic Emotions")
-                emotions_df = pd.DataFrame({
+            if "Moral Foundations Flow" in viz_options:
+                with cols[current_col]:
+                    moral_trajectories = {
+                        'care': [], 'fairness': [], 'loyalty': [],
+                        'authority': [], 'sanctity': []
+                    }
+                    
+                    for segment in segments:
+                        moral_scores = analyzer.analyze_moral_foundations(segment)
+                        for foundation in moral_trajectories.keys():
+                            moral_trajectories[foundation].append(moral_scores[foundation])
+                    
+                    moral_fig = go.Figure()
+                    colors = px.colors.qualitative.Set3[:5]
+                    
+                    for idx, (foundation, scores) in enumerate(moral_trajectories.items()):
+                        moral_fig.add_trace(go.Scatter(
+                            x=segment_labels,
+                            y=scores,
+                            name=MORAL_FOUNDATIONS[foundation],
+                            mode='lines+markers',
+                            line=dict(color=colors[idx], width=2),
+                            marker=dict(size=6, color=colors[idx])
+                        ))
+                    
+                    moral_fig.update_layout(
+                        title='Moral Foundations Flow',
+                        xaxis_title='Speech Segments',
+                        yaxis_title='Foundation Strength',
+                        yaxis=dict(range=[0, 1]),
+                        hovermode='x unified',
+                        plot_bgcolor='white',
+                        showlegend=True,
+                        height=500,
+                        legend=dict(
+                            orientation="h",
+                            yanchor="bottom",
+                            y=1.02,
+                            xanchor="right",
+                            x=1
+                        )
+                    )
+                    
+                    moral_fig.update_traces(
+                        hovertemplate="Segment: %{x}<br>Strength: %{y:.2f}<extra></extra>"
+                    )
+                    
+                    st.plotly_chart(moral_fig, use_container_width=True)
+                    if compare_mode:
+                        current_col += 1
+        
+            if "Basic Emotions Flow" in viz_options:
+                with cols[current_col]:
+                    emotions = []
+                    for segment in segments:
+                        try:
+                            result = analyzer.emotion_classifier(segment[:512])[0]
+                            emotions.append(result['label'])
+                        except:
+                            emotions.append('neutral')
+                    
+                    emotions_df = pd.DataFrame({
+                        'Segment': segment_labels,
+                        'Emotion': emotions
+                    })
+                    
+                    emotions_fig = px.bar(
+                        emotions_df,
+                        x='Segment',
+                        y='Emotion',
+                        color='Emotion',
+                        title='Basic Emotions Flow',
+                        color_discrete_sequence=px.colors.qualitative.Set2
+                    )
+                    
+                    emotions_fig.update_layout(
+                        xaxis_title='Speech Segments',
+                        yaxis_title='Emotion',
+                        showlegend=True,
+                        plot_bgcolor='white',
+                        height=500,
+                        legend=dict(
+                            orientation="h",
+                            yanchor="bottom",
+                            y=1.02,
+                            xanchor="right",
+                            x=1
+                        )
+                    )
+                    
+                    emotions_fig.update_traces(
+                        hovertemplate="Segment: %{x}<br>Emotion: %{y}<extra></extra>"
+                    )
+                    
+                    st.plotly_chart(emotions_fig, use_container_width=True)
+                    if compare_mode:
+                        current_col += 1
+        
+            # Add download button
+            if viz_options:
+                results_df = pd.DataFrame({
                     'Segment': segment_labels,
-                    'Emotion': basic_emotions
+                    'Sentiment': sentiment_scores if "Sentiment Flow" in viz_options else None,
+                    'Emotion': emotions if "Basic Emotions Flow" in viz_options else None
                 })
                 
-                emotions_fig = px.bar(
-                    emotions_df,
-                    x='Segment',
-                    y='Emotion',
-                    color='Emotion',
-                    title='Basic Emotions Flow',
-                    category_orders={'Emotion': ['joy', 'sadness', 'anger', 'fear', 'surprise', 'neutral']}
-                )
+                if "Moral Foundations Flow" in viz_options:
+                    for foundation, scores in moral_trajectories.items():
+                        results_df[f'Moral_{foundation}'] = scores
                 
-                emotions_fig.update_layout(
-                    xaxis_title='Speech Segments',
-                    yaxis_title='Emotion',
-                    showlegend=True,
-                    plot_bgcolor='white'
+                csv = results_df.to_csv(index=False)
+                st.download_button(
+                    label="Download Analysis Results",
+                    data=csv,
+                    file_name="speech_analysis_results.csv",
+                    mime="text/csv"
                 )
-                st.plotly_chart(emotions_fig)
 
         with tab3:
             status_text.text('Analyzing Linguistic Features...')