diagrams.py

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')

    ax.text(layer_count * 1.5, 1, "→ Output", va='center', fontsize=12, weight='bold')

    if memory:
        ax.text(layer_count * 0.75, 1.6, "Memory Enabled", color='purple', fontsize=10, ha='center')

    file_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    plt.savefig(file_path, bbox_inches="tight")
    plt.close()
    return file_path