Update app.py

app.py

@@ -13,42 +13,23 @@ graph = Digraph()
 graph.attr(rankdir="LR")
 graph.attr(fontsize="20")
 graph.attr(compound="true")
+graph.attr(nodesep="0.5")
 
-# Add the swim lanes
-with graph.subgraph(name="cluster_0") as c:
-    c.attr(color="gray")
-    c.attr(label=SWIM_LANES["Data Pipelines"])
-    c.node_attr.update(style="filled", color="white")
-    c.node("📊 Data Collection")
-    c.node("🧹 Data Cleaning")
-    c.node("🔧 Data Transformation")
-    
-with graph.subgraph(name="cluster_1") as c:
-    c.attr(color="gray")
-    c.attr(label=SWIM_LANES["Build and Train Models"])
-    c.node_attr.update(style="filled", color="white")
-    c.node("🔎 Feature Engineering")
-    c.node("⚙️ Model Selection")
-    c.node("🎓 Model Training")
-    
-with graph.subgraph(name="cluster_2") as c:
-    c.attr(color="gray")
-    c.attr(label=SWIM_LANES["Deploy and Predict"])
-    c.node_attr.update(style="filled", color="white")
-    c.node("🚢 Model Deployment")
-    c.node("📡 Model Serving")
-    c.node("🔮 Predictions")
-    
-# Add the RLHF step
-with graph.subgraph(name="cluster_3") as c:
-    c.attr(color="lightblue")
-    c.attr(label="Reinforcement Learning Human Feedback")
-    c.node_attr.update(style="filled", color="white")
-    c.node("👍 Feedback Collection")
-    c.node("🤔 Feedback Processing")
-    c.node("✍️ Model Updating")
-    
-# Define the edges
+# Define the nodes
+graph.node("📊 Data Collection")
+graph.node("🧹 Data Cleaning")
+graph.node("🔧 Data Transformation")
+graph.node("🔎 Feature Engineering")
+graph.node("⚙️ Model Selection")
+graph.node("🎓 Model Training")
+graph.node("🚢 Model Deployment")
+graph.node("📡 Model Serving")
+graph.node("🔮 Predictions")
+graph.node("👍 Feedback Collection")
+graph.node("🤔 Feedback Processing")
+graph.node("✍️ Model Updating")
+
+# Add the edges
 graph.edge("📊 Data Collection", "🧹 Data Cleaning")
 graph.edge("🧹 Data Cleaning", "🔧 Data Transformation")
 graph.edge("🔧 Data Transformation", "🔎 Feature Engineering")
@@ -62,5 +43,31 @@ graph.edge("👍 Feedback Collection", "🤔 Feedback Processing")
 graph.edge("🤔 Feedback Processing", "✍️ Model Updating")
 graph.edge("✍️ Model Updating", "🎓 Model Training")
 
+# Add the swim lanes
+with graph.subgraph(name="cluster_0") as c:
+    c.attr(rank="same")
+    c.attr(label=SWIM_LANES["Data Pipelines"])
+    c.edge("📊 Data Collection", "🧹 Data Cleaning", style="invis")
+    c.edge("🧹 Data Cleaning", "🔧 Data Transformation", style="invis")
+
+with graph.subgraph(name="cluster_1") as c:
+    c.attr(rank="same")
+    c.attr(label=SWIM_LANES["Build and Train Models"])
+    c.edge("🔎 Feature Engineering", "⚙️ Model Selection", style="invis")
+    c.edge("⚙️ Model Selection", "🎓 Model Training", style="invis")
+
+with graph.subgraph(name="cluster_2") as c:
+    c.attr(rank="same")
+    c.attr(label=SWIM_LANES["Deploy and Predict"])
+    c.edge("🚢 Model Deployment", "📡 Model Serving", style="invis")
+    c.edge("📡 Model Serving", "🔮 Predictions", style="invis")
+
+with graph.subgraph(name="cluster_3") as c:
+    c.attr(rank="same")
+    c.attr(label="Reinforcement Learning Human Feedback")
+    c.edge("🔮 Predictions", "👍 Feedback Collection", style="invis")
+    c.edge("👍 Feedback Collection", "🤔 Feedback Processing", style="invis")
+    c.edge("🤔 Feedback Processing", "✍️ Model Updating", style="invis")
+
 # Render the graph in Streamlit
-st.graphviz_chart(graph.source)
+st.graphviz_chart(graph.source)