Update visualization.py

visualization.py

@@ -1,82 +1,11 @@
 import gradio as gr
 import plotly.graph_objs as go
 
-def extract_data_and_explanation(text):
-    speakers_data = {}
-    current_speaker = None
-    explanation = ""
-    for line in text.split('\n'):
-        line = line.strip()
-        if line.startswith("-------"):
-            if current_speaker and explanation:
-                speakers_data[current_speaker]["explanation"] = explanation.strip()
-            explanation = ""
-            current_speaker = None
-            continue
-        if line.startswith("Speaker"):
-            current_speaker = line.strip()  # Remove the colon
-            speakers_data[current_speaker] = {}
-        elif ':' in line and current_speaker:
-            key, value = line.split(':', 1)
-            key = key.strip()
-            value = value.strip()
-            if key.lower() == "explanation":
-                explanation += value + " "
-            elif key in ["Self", "Others", "Anxiety", "Avoidance"]:
-                try:
-                    speakers_data[current_speaker][key] = float(value)
-                except ValueError:
-                    speakers_data[current_speaker][key] = value
-            else:
-                try:
-                    speakers_data[current_speaker][key] = float(value)
-                except ValueError:
-                    speakers_data[current_speaker][key] = value
-        elif line and current_speaker and not line.startswith("Explanation"):
-            explanation += line + " "
-    
-    if current_speaker and explanation:
-        speakers_data[current_speaker]["explanation"] = explanation.strip()
-
-    return speakers_data
-
-def create_bar_chart(data, title):
-    fig = go.Figure(data=[go.Bar(
-        x=list(data.keys()), 
-        y=list(data.values()),
-        text=list(data.values()),
-        textposition='auto',
-        marker_color=['red', 'green', 'blue', 'yellow', 'purple', 'orange', 'pink', 'cyan', 'magenta', 'brown'][:len(data)]
-    )])
-    fig.update_layout(title=title, xaxis_title="Traits", yaxis_title="Score")
-    fig.update_xaxes(tickangle=45)
-    return fig
-
-def create_radar_chart(data, title):
-    values = [data.get('Avoidance', 0), data.get('Self', 0), data.get('Anxiety', 0), data.get('Others', 0)]
-    fig = go.Figure(data=go.Scatterpolar(
-        r=values,
-        theta=['Avoidance', 'Self', 'Anxiety', 'Others'],
-        fill='toself'
-    ))
-    fig.update_layout(
-        polar=dict(
-            radialaxis=dict(visible=True, range=[0, max(values + [10])])
-        ),
-        showlegend=False,
-        title=title
-    )
-    return fig
-
-def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive, personalities, original_tokens, attachments_tokens, bigfive_tokens, personalities_tokens):
+def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive, personalities):
     outputs = [
         gr.update(value=status, visible=True),
         gr.update(value=exec_time, visible=True),
         gr.update(value=lang, visible=True),
-        gr.update(value=f"Original tokens: {original_tokens}", visible=True),
-        gr.update(value=f"Attachments tokens: {attachments_tokens}", visible=True),
-        gr.update(value=f"Big Five tokens: {bigfive_tokens}", visible=True),
-        gr.update(value=f"Personalities tokens: {personalities_tokens}", visible=True),
     ]
     
     for analysis_text, analysis_type in [(attachments, "Attachments"), (bigfive, "Big Five"), (personalities, "Personalities")]:
@@ -86,7 +15,7 @@ def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive,
                 if analysis_type == "Attachments":
                     chart_data = {k: v for k, v in data.items() if k in ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"] and isinstance(v, (int, float))}
                     if chart_data:
-                        fig = create_bar_chart(chart_data, f"{analysis_type} - {speaker}")
+                        fig = create_bar_chart(chart_data, f"{analysis_type} Analysis - {speaker}")
                         outputs.append(gr.update(value=fig, visible=True))
                     else:
                         outputs.append(gr.update(visible=False))
@@ -100,7 +29,7 @@ def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive,
                 else:
                     chart_data = {k: v for k, v in data.items() if k not in ["explanation"] and isinstance(v, (int, float))}
                     if chart_data:
-                        fig = create_bar_chart(chart_data, f"{analysis_type} - {speaker}")
+                        fig = create_bar_chart(chart_data, f"{analysis_type} Analysis - {speaker}")
                         outputs.append(gr.update(value=fig, visible=True))
                     else:
                         outputs.append(gr.update(visible=False))
@@ -111,8 +40,75 @@ def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive,
             else:
                 outputs.extend([gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)])
     
-    # Ensure we always return exactly 25 outputs
-    while len(outputs) < 25:
+    # Ensure we always return exactly 21 outputs
+    while len(outputs) < 21:
         outputs.append(gr.update(visible=False))
     
-    return outputs[:25]  # Trim to exactly 25 outputs
+    return outputs[:21]  # Trim to exactly 21 outputs
+
+def create_bar_chart(data, title):
+    if not data:
+        return None
+    fig = go.Figure(data=[go.Bar(
+        x=list(data.keys()), 
+        y=list(data.values()),
+        text=list(data.values()),
+        textposition='auto',
+        marker_color=['red', 'green', 'blue', 'yellow', 'purple', 'orange', 'pink', 'cyan', 'magenta', 'brown'][:len(data)]
+    )])
+    fig.update_layout(title=title, xaxis_title="Traits", yaxis_title="Score")
+    fig.update_xaxes(tickangle=45)
+    return fig
+
+def create_radar_chart(data, title):
+    if not data:
+        return None
+    values = [data.get('Avoidance', 0), data.get('Self', 0), data.get('Anxiety', 0), data.get('Others', 0)]
+    fig = go.Figure(data=go.Scatterpolar(
+        r=values,
+        theta=['Avoidance', 'Self', 'Anxiety', 'Others'],
+        fill='toself'
+    ))
+    fig.update_layout(
+        polar=dict(
+            radialaxis=dict(visible=True, range=[0, max(values + [10])])
+        ),
+        showlegend=False,
+        title=title
+    )
+    return fig
+
+# Add this function to extract data from the text
+def extract_data_and_explanation(text):
+    speakers_data = {}
+    current_speaker = None
+    explanation = ""
+    for line in text.split('\n'):
+        line = line.strip()
+        if line.startswith("-----------------------"):
+            if current_speaker and explanation:
+                speakers_data[current_speaker]["explanation"] = explanation.strip()
+            explanation = ""
+            current_speaker = None
+            continue
+        if line.startswith("Speaker"):
+            current_speaker = line.strip()
+            speakers_data[current_speaker] = {}
+        elif ':' in line and current_speaker:
+            key, value = line.split(':', 1)
+            key = key.strip()
+            value = value.strip()
+            if key.lower() == "explanation":
+                explanation += value + " "
+            else:
+                try:
+                    speakers_data[current_speaker][key] = float(value)
+                except ValueError:
+                    speakers_data[current_speaker][key] = value
+        elif line and current_speaker and not line.startswith("Explanation"):
+            explanation += line + " "
+    
+    if current_speaker and explanation:
+        speakers_data[current_speaker]["explanation"] = explanation.strip()
+
+    return speakers_data