Update diagrams.py

diagrams.py

@@ -1,23 +1,28 @@
 import matplotlib.pyplot as plt
-import tempfile
 
-def get_transformer_diagram(config):
-    fig, ax = plt.subplots(figsize=(8, 2))
-    ax.axis("off")
-    
-    layer_count = config["layers"]
-    memory = config.get("memory", False)
-    
-    for i in range(layer_count):
-        ax.add_patch(plt.Rectangle((i * 1.5, 0.5), 1.2, 1, edgecolor='black', facecolor='skyblue'))
-        ax.text(i * 1.5 + 0.6, 1, f"Layer {i+1}", ha='center', va='center')
+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)
 
-    ax.text(layer_count * 1.5, 1, "→ Output", va='center', fontsize=12, weight='bold')
+        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)
 
-    if memory:
-        ax.text(layer_count * 0.75, 1.6, "Memory Enabled", color='purple', fontsize=10, ha='center')
+        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)
 
-    file_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
-    plt.savefig(file_path, bbox_inches="tight")
-    plt.close()
-    return file_path
+        ax.set_ylim(0.9, 1.15)
+        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)