Switch to Gradio version

app.py

@@ -2,8 +2,8 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
 
-def plot_3d(max_t, num_points):
-    # Generate sequence
+def plot_3d(max_t, num_points, sentence, show_words):
+    # Generate spiral
     t = np.linspace(0, max_t, num_points)
     x = np.sin(t)
     y = np.cos(t)
@@ -15,19 +15,46 @@ def plot_3d(max_t, num_points):
     ax.plot(x, y, z, label="3D Spiral")
     ax.legend()
 
+    # Add token labels if requested
+    if show_words and sentence.strip():
+        tokens = sentence.strip().split()
+        idxs = np.linspace(0, len(t) - 1, len(tokens), dtype=int)
+        for i, token in zip(idxs, tokens):
+            ax.text(x[i], y[i], z[i], token, fontsize=9, color="red")
+
     return fig
 
-# Build Gradio interface with sliders
-demo = gr.Interface(
-    fn=plot_3d,
-    inputs=[
-        gr.Slider(5, 50, value=20, step=1, label="Spiral Length"),
-        gr.Slider(100, 2000, value=500, step=50, label="Number of Points")
-    ],
-    outputs="plot",
-    title="3D Hidden States Visualization",
-    description="Adjust the spiral length and number of points to explore the 3D curve."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# 3D Hidden States Visualization")
+    gr.Markdown(
+        """
+        This plot shows a **3D spiral**, generated using sine and cosine for the x/y axes 
+        and a linear sequence for the z-axis:
+
+        - $x = \\sin(t)$  
+        - $y = \\cos(t)$  
+        - $z = t$  
+
+        Together, these equations trace a spiral around the z-axis.  
+        Think of it as an analogy for **hidden states in a neural network**, 
+        which evolve over time (the z-axis), while oscillating in complex patterns (x & y axes).  
+
+        ✨ Try typing your own sentence — each word will be placed along the spiral, 
+        showing how tokens could be mapped into hidden state space.
+        """
+    )
+
+    with gr.Row():
+        max_t = gr.Slider(5, 50, value=20, step=1, label="Spiral Length")
+        num_points = gr.Slider(100, 2000, value=500, step=50, label="Number of Points")
+
+    sentence = gr.Textbox(value="I love hidden states in transformers", label="Sentence")
+    show_words = gr.Checkbox(label="Show Tokens", value=True)
+
+    plot = gr.Plot()
+
+    btn = gr.Button("Generate Plot")
+    btn.click(plot_3d, inputs=[max_t, num_points, sentence, show_words], outputs=plot)
 
 if __name__ == "__main__":
     demo.launch()