Update app.py

app.py

@@ -4,7 +4,38 @@ import numpy as np
 from transformers import pipeline, RobertaForSequenceClassification, RobertaTokenizer
 from motif_tagging import detect_motifs
 import re
-
+import matplotlib.pyplot as plt
+import io
+from PIL import Image
+from datetime import datetime
+
+# --- Timeline Visualization Function ---
+def generate_risk_stage_timeline(dates, stages):
+    stage_labels = {
+        1: "Tension",
+        2: "Escalation",
+        3: "Reconciliation",
+        4: "Calm"
+    }
+    try:
+        parsed_dates = [datetime.strptime(d, "%Y-%m-%d") for d in dates]
+    except Exception:
+        parsed_dates = list(range(len(dates)))
+
+    fig, ax = plt.subplots(figsize=(6, 2.5))
+    ax.step(parsed_dates, stages, where='mid', color='purple', linewidth=2, marker='o')
+    ax.set_yticks(list(stage_labels.keys()))
+    ax.set_yticklabels([stage_labels[s] for s in stage_labels])
+    ax.set_title("Risk Stage Timeline")
+    ax.set_xlabel("Date")
+    ax.set_ylabel("Risk Stage")
+    ax.grid(True)
+    plt.tight_layout()
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    return Image.open(buf)
 # --- SST Sentiment Model ---
 sst_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
 
@@ -300,7 +331,8 @@ def analyze_composite(msg1, date1, msg2, date2, msg3, date3, *answers_and_none):
     out += f"\n\n{stage_text}"
     out += darvo_blurb
 
-    return out
+    timeline_image = generate_risk_stage_timeline(dates_used, stages)
+    return out, timeline_image
     
 message_date_pairs = [(gr.Textbox(label=f"Message {i+1}"), gr.Textbox(label=f"Date {i+1} (optional)", placeholder="e.g. 2025-04-22")) for i in range(3)]
 textbox_inputs = [item for pair in message_date_pairs for item in pair]  # Flatten for Gradio input
@@ -310,7 +342,10 @@ none_box = gr.Checkbox(label="None of the above")
 iface = gr.Interface(
     fn=analyze_composite,
     inputs=textbox_inputs + quiz_boxes + [none_box],
-    outputs=gr.Textbox(label="Results"),
+    outputs=[
+        gr.Textbox(label="Results"),
+        gr.Image(label="Risk Stage Timeline", type="pil")
+    ],
     title="Abuse Pattern Detector + Escalation Quiz",
     allow_flagging="manual"
 )