Update diagram.py

diagram.py

@@ -1,28 +1,14 @@
-import matplotlib.pyplot as plt
-
-def plot_architecture_diagram(config):
-    try:
-        fig, ax = plt.subplots(figsize=(6, 2))  # reduced size
-        layers = config.get("layers", 4)
-        heads = config.get("attention_heads", 4)
-        ffn = config.get("ffn_dim", 1024)
-        mem = config.get("memory", False)
+# diagram.py
 
-        labels = []
-        for i in range(layers):
-            lbl = f"L{i+1}\n{heads}H\n{ffn}F"
-            if mem and i == layers // 2:
-                lbl += "\nMem"
-            labels.append(lbl)
-
-        ax.plot(range(layers), [1] * layers, "o-", linewidth=2)
-        for i, label in enumerate(labels):
-            ax.text(i, 1.02, label, ha="center", fontsize=8)
+import matplotlib.pyplot as plt
+import tempfile
 
-        ax.set_ylim(0.9, 1.15)
-        ax.axis("off")
+def show_transformer_diagram():
+    fig, ax = plt.subplots(figsize=(6, 4))
+    ax.text(0.5, 0.6, 'Input → Embedding → Self-Attention → FFN → Output',
+            fontsize=12, ha='center', va='center', wrap=True)
+    ax.axis('off')
 
-        fig.savefig("architecture_diagram.png", bbox_inches="tight", dpi=150)
-        plt.close(fig)
-    except Exception as e:
-        print("⚠️ Diagram plot error:", e)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
+        plt.savefig(tmpfile.name, bbox_inches="tight")
+        return tmpfile.name