Update app.py

app.py

@@ -8,18 +8,24 @@ from huggingface_hub import HfApi
 import plotly.graph_objects as go
 import time
 from datetime import datetime
+import threading
 
 # Cyberpunk and Loading Animation Styling
 def setup_cyberpunk_style():
     st.markdown("""
         <style>
-        @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&display=swap');
-        @import url('https://fonts.googleapis.com/css2?family=Share+Tech+Mono&display=swap');
-        
+        body, button, input, select, textarea {
+            font-family: 'Orbitron', sans-serif !important;
+            color: #00ff9d !important;
+        }
         .stApp {
             background: radial-gradient(circle, rgba(0, 0, 0, 0.95) 20%, rgba(0, 50, 80, 0.95) 90%);
             color: #00ff9d;
             font-family: 'Orbitron', sans-serif;
+            font-size: 16px;
+            line-height: 1.6;
+            padding: 20px;
+            box-sizing: border-box;
         }
         
         .main-title {
@@ -145,7 +151,9 @@ def initialize_model(model_name="gpt2"):
 # Load Dataset Function with Uploaded File Option
 def load_dataset(data_source="demo", tokenizer=None, uploaded_file=None):
     if data_source == "demo":
-        data = ["Sample text data for model training. This can be replaced with actual data for better performance."]
+        data = ["In the neon-lit streets of Neo-Tokyo, a lone hacker fights against the oppressive megacorporations.", 
+                "The rain falls in sheets, washing away the bloodstains from the alleyways.", 
+                "She plugs into the matrix, seeking answers to questions that have haunted her for years."]
     elif uploaded_file is not None:
         if uploaded_file.name.endswith(".txt"):
             data = [uploaded_file.read().decode("utf-8")]
@@ -160,7 +168,7 @@ def load_dataset(data_source="demo", tokenizer=None, uploaded_file=None):
     return dataset
 
 # Train Model Function with Customized Progress Bar
-def train_model(model, train_dataset, tokenizer, epochs=3, batch_size=4):
+def train_model(model, train_dataset, tokenizer, epochs=3, batch_size=4, progress_callback=None):
     training_args = TrainingArguments(
         output_dir="./results",
         overwrite_output_dir=True,
@@ -179,14 +187,26 @@ def train_model(model, train_dataset, tokenizer, epochs=3, batch_size=4):
         args=training_args,
         data_collator=data_collator,
         train_dataset=train_dataset,
+        callbacks=[ProgressCallback(progress_callback)]
     )
 
     trainer.train()
 
+class ProgressCallback(TrainerCallback):
+    def __init__(self, progress_callback):
+        super().__init__()
+        self.progress_callback = progress_callback
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        loss = state.log_history[-1]['loss']
+        generation = state.global_step // args.gradient_accumulation_steps + 1
+        individual = args.gradient_accumulation_steps
+        self.progress_callback(loss, generation, individual)
+
 # Main App Logic
 def main():
     setup_cyberpunk_style()
-    st.markdown('<h1 class="main-title">Cyberpunk Neural Training Hub</h1>', unsafe_allow_html=True)
+    st.markdown('<h1 class="main-title">Neural Training Hub</h1>', unsafe_allow_html=True)
     
     # Initialize model and tokenizer
     model, tokenizer = initialize_model()
@@ -225,6 +245,15 @@ def main():
     # Load Dataset
     train_dataset = load_dataset(data_source, tokenizer, uploaded_file=uploaded_file)
 
+    # Chatbot Interaction
+    if st.checkbox("Enable Chatbot"):
+        user_input = st.text_input("You:", placeholder="Type your message here...")
+        if user_input:
+            inputs = tokenizer(user_input, return_tensors="pt")
+            outputs = model.generate(inputs['input_ids'], max_length=100, num_return_sequences=1)
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            st.write("Bot:", response)
+
     # Go Button to Start Training
     if st.button("Go"):
         progress_placeholder = st.empty()
@@ -233,22 +262,21 @@ def main():
 
         dashboard = TrainingDashboard()
 
-        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
+        def train_progress(loss, generation, individual):
+            progress = (generation + 1) / dashboard.metrics['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)
+            dashboard.update(loss=loss, generation=generation, individual=individual)
+
+        thread = threading.Thread(target=train_model, args=(model, train_dataset, tokenizer, training_epochs, batch_size, train_progress))
+        thread.start()
+        loading_animation.markdown("""
+            <div class="loading-animation"></div>
+        """, unsafe_allow_html=True)
+        thread.join()
         
         loading_animation.empty()
         st.success("Training Complete!")