Update visualization.py

visualization.py

@@ -6,6 +6,7 @@ from collections import Counter
 def extract_speaker_data(text):
     speakers = {}
     current_speaker = None
+    explanation = ""
     for line in text.split('\n'):
         if line.lower().startswith("speaker"):
             current_speaker = line.split(":")[1].strip() if ":" in line else line.split()[1]
@@ -18,7 +19,9 @@ def extract_speaker_data(text):
                 match = re.search(r"[-+]?\d*\.\d+|\d+", value)
                 if match:
                     speakers[current_speaker][key.strip()] = float(match.group())
-    return speakers
+        elif line.lower().startswith("explanation:"):
+            explanation = line.split("Explanation:", 1)[1].strip()
+    return speakers, explanation
 
 def create_bar_chart(data, title, speaker):
     fig = go.Figure(data=[go.Bar(
@@ -30,12 +33,17 @@ def create_bar_chart(data, title, speaker):
     return fig
 
 def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive, personalities):
-    charts = []
+    outputs = [
+        gr.update(value=status, visible=True),
+        gr.update(value=exec_time, visible=True),
+        gr.update(value=lang, visible=True),
+    ]
     
-    for analysis_text in [attachments, bigfive, personalities]:
-        speakers_data = extract_speaker_data(analysis_text)
+    for analysis_text, analysis_type in [(attachments, "Attachments"), (bigfive, "Big Five"), (personalities, "Personalities")]:
+        speakers_data, explanation = extract_speaker_data(analysis_text)
         if not speakers_data:
             print(f"No speaker data extracted from: {analysis_text}")
+            continue
 
         # Determine the two main speakers
         speaker_counts = Counter(speakers_data.keys())
@@ -43,20 +51,20 @@ def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive,
         
         for speaker in main_speakers:
             data = speakers_data.get(speaker, {})
-            attachment_data = {k: v for k, v in data.items() if k in ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"]}
-            if attachment_data:
-                charts.append(create_bar_chart(attachment_data, "Attachment Styles", speaker))
-            
-            bigfive_data = {k: v for k, v in data.items() if k in ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"]}
-            if bigfive_data:
-                charts.append(create_bar_chart(bigfive_data, "Big Five Traits", speaker))
+            if analysis_type == "Attachments":
+                chart_data = {k: v for k, v in data.items() if k in ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"]}
+                if chart_data:
+                    outputs.append(gr.Plot(create_bar_chart(chart_data, "Attachment Styles", speaker)))
+            elif analysis_type == "Big Five":
+                chart_data = {k: v for k, v in data.items() if k in ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"]}
+                if chart_data:
+                    outputs.append(gr.Plot(create_bar_chart(chart_data, "Big Five Traits", speaker)))
+            elif analysis_type == "Personalities":
+                chart_data = {k: v for k, v in data.items() if k in ["Depressed", "Paranoid", "Schizoid-Schizotypal", "Antisocial-Psychopathic", "Borderline-Dysregulated", "Hysteric-Histrionic", "Narcissistic", "Anxious-Avoidant", "Dependent-Victimized", "Obsessional"]}
+                if chart_data:
+                    outputs.append(gr.Plot(create_bar_chart(chart_data, "Personality Traits", speaker)))
             
-            personality_data = {k: v for k, v in data.items() if k in ["Depressed", "Paranoid", "Schizoid-Schizotypal", "Antisocial-Psychopathic", "Borderline-Dysregulated", "Hysteric-Histrionic", "Narcissistic", "Anxious-Avoidant", "Dependent-Victimized", "Obsessional"]}
-            if personality_data:
-                charts.append(create_bar_chart(personality_data, "Personality Traits", speaker))
+            # Add explanation box after each graph
+            outputs.append(gr.Textbox(value=explanation, label=f"{analysis_type} Explanation - Speaker {speaker}"))
 
-    return [
-        gr.update(value=status, visible=True),
-        gr.update(value=exec_time, visible=True),
-        gr.update(value=lang, visible=True),
-    ] + charts
+    return outputs