Update app.py

app.py

@@ -1,48 +1,80 @@
+# Imports
 import streamlit as st
 import numpy as np
 import torch
+import random
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
 from datasets import Dataset
+from huggingface_hub import HfApi
+import plotly.graph_objects as go
 import time
 from datetime import datetime
-import plotly.graph_objects as go
-from huggingface_hub import HfApi, HfFolder
-
-# Initialize Hugging Face Authentication
-def huggingface_login():
-    token = st.text_input("Hugging Face Token", type="password")
-    if token:
-        HfFolder.save_token(token)
-        api = HfApi()
-        user_info = api.whoami(token)
-        st.sidebar.write(f"Logged in as: {user_info['name']}")
-        return token
-    else:
-        st.warning("Please enter your Hugging Face token")
-        return None
 
-# Advanced Cyberpunk Styling
-def setup_advanced_cyberpunk_style():
+# 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');
-        .main-title { font-family: 'Orbitron', sans-serif; font-size: 40px; color: #00ffea; }
-        /* Additional CSS styling for dashboard, progress bar, and background */
-        </style>
-    """, unsafe_allow_html=True)
+        
+        .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;
+        }
+        
+        .main-title {
+            text-align: center;
+            font-size: 4em;
+            color: #00ff9d;
+            letter-spacing: 4px;
+            animation: glow 2s ease-in-out infinite alternate;
+        }
+        
+        @keyframes glow {
+            from {text-shadow: 0 0 5px #00ff9d, 0 0 10px #00ff9d;}
+            to {text-shadow: 0 0 15px #00b8ff, 0 0 20px #00b8ff;}
+        }
 
-# Initialize Model and Tokenizer
-def initialize_model():
-    model = GPT2LMHeadModel.from_pretrained("gpt2")
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        .stButton > 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;
+        }
+        
+        .stButton > button:hover {
+            transform: scale(1.1);
+            box-shadow: 0 0 20px rgba(0, 255, 157, 0.5);
+        }
 
-    # Set padding token to eos_token
-    tokenizer.pad_token = tokenizer.eos_token
-    return model, tokenizer
+        .progress-bar-container {
+            background: rgba(0, 0, 0, 0.5);
+            border-radius: 15px;
+            overflow: hidden;
+            width: 100%;
+            height: 30px;
+            position: relative;
+            margin: 10px 0;
+        }
+        
+        .progress-bar {
+            height: 100%;
+            width: 0%;
+            background: linear-gradient(45deg, #00ff9d, #00b8ff);
+            transition: width 0.5s ease;
+        }
+        
+        </style>
+    """, unsafe_allow_html=True)
 
-# Prepare Dataset
+# Prepare Dataset Function with Padding Token Fix
 def prepare_dataset(data, tokenizer, block_size=128):
+    tokenizer.pad_token = tokenizer.eos_token
     def tokenize_function(examples):
         return tokenizer(examples['text'], truncation=True, max_length=block_size, padding='max_length')
 
@@ -52,145 +84,114 @@ def prepare_dataset(data, tokenizer, block_size=128):
     tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
     return tokenized_dataset
 
-# Training Dashboard Class
+# Training Dashboard Class with Enhanced Display
 class TrainingDashboard:
     def __init__(self):
         self.metrics = {
             'current_loss': 0,
             'best_loss': float('inf'),
             'generation': 0,
+            'individual': 0,
             'start_time': time.time(),
             'training_speed': 0
         }
         self.history = []
-        
-    def update(self, loss, generation):
+
+    def update(self, loss, generation, individual):
         self.metrics['current_loss'] = loss
         self.metrics['generation'] = generation
+        self.metrics['individual'] = individual
         if loss < self.metrics['best_loss']:
             self.metrics['best_loss'] = loss
+        
         elapsed_time = time.time() - self.metrics['start_time']
-        self.metrics['training_speed'] = generation / elapsed_time
+        self.metrics['training_speed'] = (generation * individual) / elapsed_time
         self.history.append({'loss': loss, 'timestamp': datetime.now().strftime('%H:%M:%S')})
-    
-    def display(self):
-        st.write(f"**Generation:** {self.metrics['generation']}")
-        st.write(f"**Current Loss:** {self.metrics['current_loss']:.4f}")
-        st.write(f"**Best Loss:** {self.metrics['best_loss']:.4f}")
-        st.write(f"**Training Speed:** {self.metrics['training_speed']:.2f} generations/sec")
-
-# Display Progress Bar
-def display_progress(progress):
-    st.markdown(f"""
-        <div class="progress-bar-container">
-            <div class="progress-bar" style="width: {progress * 100}%"></div>
-        </div>
-    """, unsafe_allow_html=True)
 
-# Custom Genetic Algorithm
-class GeneticAlgorithm:
-    def __init__(self, model, tokenizer, dataset, population_size, mutation_rate=0.1):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.dataset = dataset
-        self.population_size = population_size
-        self.mutation_rate = mutation_rate
-        self.population = [self.clone_model() for _ in range(population_size)]
-    
-    def clone_model(self):
-        # Create a clone of the model
-        return GPT2LMHeadModel.from_pretrained("gpt2")
-
-    def evaluate_fitness(self, model):
-        # Calculate the loss for a given model on the dataset
-        trainer = Trainer(
-            model=model,
-            args=TrainingArguments(output_dir="./results", per_device_train_batch_size=2, num_train_epochs=1),
-            train_dataset=self.dataset,
-            data_collator=DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False),
-        )
-        train_result = trainer.train()
-        return train_result.training_loss
-
-    def select_best_models(self, num_best=2):
-        # Selects the top models based on fitness (loss)
-        fitness_scores = [(self.evaluate_fitness(model), model) for model in self.population]
-        fitness_scores.sort(key=lambda x: x[0])  # Sort by loss
-        best_models = [model for _, model in fitness_scores[:num_best]]
-        return best_models
-
-    def crossover(self, parent1, parent2):
-        # Perform crossover by combining layers from both parents
-        child = self.clone_model()
-        for (child_param, param1, param2) in zip(child.parameters(), parent1.parameters(), parent2.parameters()):
-            # Randomly choose parameters from each parent based on crossover probability
-            if np.random.rand() > 0.5:
-                child_param.data = param1.data.clone()
-            else:
-                child_param.data = param2.data.clone()
-        return child
-
-    def mutate(self, model):
-        # Mutate model by slightly adjusting its weights
-        for param in model.parameters():
-            if np.random.rand() < self.mutation_rate:
-                mutation_tensor = torch.randn_like(param) * 0.02
-                param.data += mutation_tensor
-
-    def generate_new_population(self):
-        best_models = self.select_best_models()
-        new_population = []
-        while len(new_population) < self.population_size:
-            parent1, parent2 = np.random.choice(best_models, 2, replace=False)
-            child = self.crossover(parent1, parent2)
-            self.mutate(child)
-            new_population.append(child)
-        self.population = new_population
-
-# Training Loop with Genetic Algorithm and Loading Screen
-def training_loop(dashboard, ga, num_generations):
-    with st.spinner("Training in progress..."):
-        for generation in range(1, num_generations + 1):
-            best_loss = min([ga.evaluate_fitness(model) for model in ga.population])
-            dashboard.update(best_loss, generation)
-            progress = generation / num_generations
-            display_progress(progress)
-            dashboard.display()
-            ga.generate_new_population()
-            time.sleep(0.5)  # Simulate delay for each generation
-
-# Main Function
-def main():
-    setup_advanced_cyberpunk_style()
-    st.markdown('<h1 class="main-title">Neural Evolution GPT-2 Training Hub</h1>', unsafe_allow_html=True)
+# Define Model Initialization
+def initialize_model(model_name="gpt2"):
+    model = GPT2LMHeadModel.from_pretrained(model_name)
+    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+    return model, tokenizer
 
-    # Hugging Face Account Login
-    token = huggingface_login()
-    if token is None:
-        return
+# Load Dataset Function
+def load_dataset(data_source="demo", tokenizer=None):
+    if data_source == "demo":
+        data = ["Sample text data for model training. This can be replaced with actual data for better performance."]
+    else:
+        data = ["Loaded data from uploaded text file."]
+    dataset = prepare_dataset(data, tokenizer)
+    return dataset
+
+# Train Model Function with Customized Progress Bar
+def train_model(model, train_dataset, tokenizer, epochs=3, batch_size=4):
+    training_args = TrainingArguments(
+        output_dir="./results",
+        overwrite_output_dir=True,
+        num_train_epochs=epochs,
+        per_device_train_batch_size=batch_size,
+        save_steps=10_000,
+        save_total_limit=2,
+        logging_dir="./logs",
+        logging_steps=100,
+    )
+    
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-    # Load Model and Tokenizer
-    model, tokenizer = initialize_model()
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+    )
 
-    # Prepare Data
-    data = ["Sample training text"] * 10  # Replace with real data
-    train_dataset = prepare_dataset(data, tokenizer)
+    trainer.train()
 
-    # Initialize Dashboard
-    dashboard = TrainingDashboard()
-    
-    # Sidebar Configuration
-    st.sidebar.markdown("### Training Parameters")
-    num_generations = st.sidebar.slider("Generations", 1, 50, 10)
-    population_size = st.sidebar.slider("Population Size", 4, 20, 10)
-    mutation_rate = st.sidebar.slider("Mutation Rate", 0.01, 0.5, 0.1)
+# Main App Logic
+def main():
+    setup_cyberpunk_style()
+    st.markdown('<h1 class="main-title">Cyberpunk Neural Training Hub</h1>', unsafe_allow_html=True)
     
-    # Initialize Genetic Algorithm
-    ga = GeneticAlgorithm(model, tokenizer, train_dataset, population_size, mutation_rate)
+    # Initialize model and tokenizer
+    model, tokenizer = initialize_model()
+
+    # Sidebar Configuration with Additional Options
+    with st.sidebar:
+        st.markdown("### Configuration Panel")
+        training_epochs = st.slider("Training Epochs", min_value=1, max_value=5, value=3)
+        batch_size = st.slider("Batch Size", min_value=2, max_value=8, value=4)
+        model_choice = st.selectbox("Model Selection", ("gpt2", "distilgpt2", "gpt2-medium"))
+        data_source = st.selectbox("Data Source", ("demo", "uploaded file"))
+        custom_learning_rate = st.slider("Learning Rate", min_value=1e-6, max_value=5e-4, value=3e-5, step=1e-6)
+        
+        advanced_toggle = st.checkbox("Advanced Training Settings")
+        if advanced_toggle:
+            warmup_steps = st.slider("Warmup Steps", min_value=0, max_value=500, value=100)
+            weight_decay = st.slider("Weight Decay", min_value=0.0, max_value=0.1, step=0.01, value=0.01)
+        else:
+            warmup_steps = 100
+            weight_decay = 0.01
+
+    # Load Dataset
+    train_dataset = load_dataset(data_source, tokenizer)
+
+    # Start Training with Progress Bar
+    progress_placeholder = st.empty()
+    st.markdown("### Model Training Progress")
+
+    for epoch in range(training_epochs):
+        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)
     
-    # Run Training
-    if st.button("Start Training"):
-        training_loop(dashboard, ga, num_generations)
+    st.success("Training Complete!")
 
 if __name__ == "__main__":
     main()