Update app.py

app.py

@@ -1,4 +1,3 @@
-# Imports
 import streamlit as st
 import numpy as np
 import torch
@@ -67,6 +66,36 @@ def setup_cyberpunk_style():
             transition: width 0.5s ease;
         }
         
+        .go-button {
+            font-family: 'Orbitron', sans-serif;
+            background: linear-gradient(45deg, #00ff9d, #00b8ff);
+            color: #000;
+            font-size: 1.1em;
+            padding: 10px 20px;
+            border: none;
+            border-radius: 8px;
+            transition: all 0.3s ease;
+            cursor: pointer;
+        }
+        
+        .go-button:hover {
+            transform: scale(1.1);
+            box-shadow: 0 0 20px rgba(0, 255, 157, 0.5);
+        }
+        
+        .loading-animation {
+            display: inline-block;
+            width: 20px;
+            height: 20px;
+            border: 3px solid #00ff9d;
+            border-radius: 50%;
+            border-top-color: transparent;
+            animation: spin 1s ease-in-out infinite;
+        }
+        
+        @keyframes spin {
+            to {transform: rotate(360deg);}
+        }
         </style>
     """, unsafe_allow_html=True)
 
@@ -196,23 +225,35 @@ def main():
     # Load Dataset
     train_dataset = load_dataset(data_source, tokenizer, uploaded_file=uploaded_file)
 
-    # Start Training with Progress Bar
-    progress_placeholder = st.empty()
-    st.markdown("### Model Training Progress")
+    # Go Button to Start Training
+    if st.button("Go"):
+        progress_placeholder = st.empty()
+        loading_animation = st.empty()
+        st.markdown("### Model Training Progress")
+
+        dashboard = TrainingDashboard()
 
-    for epoch in range(training_epochs):
-        train_model(model, train_dataset, tokenizer, epochs=1, batch_size=batch_size)
+        for epoch in range(training_epochs):
+            loading_animation.markdown("""
+                <div class="loading-animation"></div>
+            """, unsafe_allow_html=True)
+            
+            train_model(model, train_dataset, tokenizer, epochs=1, batch_size=batch_size)
+            
+            # Update Progress Bar
+            progress = (epoch + 1) / training_epochs * 100
+            progress_placeholder.markdown(f"""
+                <div class="progress-bar-container">
+                    <div class="progress-bar" style="width: {progress}%;"></div>
+                </div>
+            """, unsafe_allow_html=True)
+            
+            dashboard.update(loss=0, generation=epoch + 1, individual=batch_size)
         
-        # Update Progress Bar
-        progress = (epoch + 1) / training_epochs * 100
-        progress_placeholder.markdown(f"""
-            <div class="progress-bar-container">
-                <div class="progress-bar" style="width: {progress}%;"></div>
-            </div>
-        """, unsafe_allow_html=True)
-    
-    st.success("Training Complete!")
+        loading_animation.empty()
+        st.success("Training Complete!")
+        st.write("Training Metrics:")
+        st.write(dashboard.metrics)
 
 if __name__ == "__main__":
-    main()
-
+    main()