Update app.py

app.py

@@ -1,26 +1,87 @@
+# app.py  (Hugging Face Space - runs GPT-2 + TransformerLens locally)
+import os
+import json
 import streamlit as st
 import requests
+from datetime import datetime
 
-API_URL = "https://activation-patching-api.onrender.com"
+# Import your model utilities (make sure model_utils.py is in the same repo on HF)
+from model_utils import generate_text, run_activation_patching  # uploaded at /mnt/data/model_utils.py
 
-st.title("Mechanistic Analysis Interface")
-st.write("Run GPT-2 generation + activation patching experiments.")
+# Configure Render URL (set your actual Render URL here)
+# Example: https://activation-patching-api.onrender.com
+RENDER_API_BASE = st.secrets.get("render_url") or st.sidebar.text_input(
+    "Render API base URL", value="https://activation-patching-api.onrender.com"
+)
+SAVE_ENDPOINT = RENDER_API_BASE.rstrip("/") + "/save"
 
-prompt = st.text_area("Enter your sentence:")
+st.title("Mechanistic Analysis Interface (HF Space)")
+st.write("This Streamlit app runs GPT-2 + activation patching locally, then saves metadata to your Render backend.")
+
+prompt = st.text_area("Enter your sentence / prompt", height=150)
+max_length = st.sidebar.slider("Max generation length", 20, 200, 60)
 
 if st.button("Run Experiment"):
-    if prompt:
-        with st.spinner("Running experiment..."):
-            response = requests.post(f"{API_URL}/generate", json={"prompt": prompt})
-            data = response.json()
-        
+    if not prompt.strip():
+        st.warning("Please enter a prompt.")
+    else:
+        with st.spinner("Generating text with GPT-2..."):
+            try:
+                generated = generate_text(prompt, max_length=max_length)
+            except Exception as e:
+                st.error(f"Error running generation: {e}")
+                raise
+
         st.subheader("Generated Text")
-        st.write(data["generated_text"])
-        
-        st.subheader("Activation Patching Traces")
-        st.write(data["activations"])
-        
+        st.write(generated)
+
+        with st.spinner("Running activation patching (TransformerLens)..."):
+            try:
+                activations = run_activation_patching(prompt)
+            except Exception as e:
+                st.error(f"Error running activation patching: {e}")
+                raise
+
+        st.subheader("Activation traces (sample)")
+        # activations can be large — summarize for display
+        sample = {k: (v.shape if hasattr(v, "shape") else type(v).__name__) for k, v in list(activations.items())[:10]}
+        st.json(sample)
+
+        # Ask explanation agent or placeholder — for now, simple summary:
+        explanation = "Explanation placeholder: top influencing layers ... (expand with LangGraph later)"
+
         st.subheader("Explanation")
-        st.write(data["explanation"])
-        
-        st.success(f"Experiment saved with ID: {data['id']}")
+        st.write(explanation)
+
+        # Prepare data to save to Render
+        payload = {
+            "prompt": prompt,
+            "generated_text": generated,
+            # Convert activations to a compact serializable form: keep shapes and optionally min/max
+            "activation_traces": json.dumps({
+                k: {
+                    "shape": getattr(v, "shape", None),
+                    "min": float(v.min()) if hasattr(v, "min") else None,
+                    "max": float(v.max()) if hasattr(v, "max") else None
+                } for k, v in activations.items()
+            }),
+            "explanation": explanation
+        }
+
+        # Save to Render
+        try:
+            res = requests.post(SAVE_ENDPOINT, json=payload, timeout=30)
+            st.write("Save status:", res.status_code)
+            st.write("Save response:", res.text)
+            if res.ok:
+                data = res.json()
+                st.success(f"Experiment saved with ID {data.get('id')}")
+            else:
+                st.error(f"Failed to save experiment: {res.text}")
+        except Exception as e:
+            st.error(f"Error saving to Render: {e}")
+
+st.markdown("---")
+st.write("Notes:")
+st.write("- This app runs heavy ML locally in the HF Space container; Render is used only to persist metadata.")
+st.write("- If you want LangGraph explanation, we can call a low-cost open model here or run agents locally.")