Update app.py

app.py

@@ -6,6 +6,20 @@ from datasets import Dataset
 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():
@@ -13,7 +27,8 @@ def setup_advanced_cyberpunk_style():
         <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');
-        /* Additional styling as provided previously */
+        .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)
 
@@ -30,10 +45,7 @@ def prepare_dataset(data, tokenizer, block_size=128):
 
     raw_dataset = Dataset.from_dict({'text': data})
     tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=['text'])
-    tokenized_dataset = tokenized_dataset.map(
-        lambda examples: {'labels': examples['input_ids']},
-        batched=True
-    )
+    tokenized_dataset = tokenized_dataset.map(lambda examples: {'labels': examples['input_ids']}, batched=True)
     tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
     return tokenized_dataset
 
@@ -72,36 +84,88 @@ def display_progress(progress):
         </div>
     """, unsafe_allow_html=True)
 
-# Fitness Calculation (Placeholder for actual loss computation)
-def compute_loss(model, dataset):
-    # Placeholder for real loss computation with Trainer API or custom logic
-    trainer = Trainer(
-        model=model,
-        args=TrainingArguments(output_dir="./results", per_device_train_batch_size=2, num_train_epochs=1),
-        train_dataset=dataset,
-        data_collator=DataCollatorForLanguageModeling(tokenizer=model.config._name_or_path, mlm=False),
-    )
-    train_result = trainer.train()
-    return train_result.training_loss
-
-# Training Loop with Loading Screen
-def training_loop(dashboard, model, dataset, num_generations, population_size):
+# 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):
-            # Simulated population loop
-            for individual in range(population_size):
-                loss = compute_loss(model, dataset)
-                dashboard.update(loss, generation)
-                progress = generation / num_generations
-                display_progress(progress)
-                dashboard.display()
-                time.sleep(1)  # Simulate delay for each individual training
+            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)
 
+    # Hugging Face Account Login
+    token = huggingface_login()
+    if token is None:
+        return
+
     # Load Model and Tokenizer
     model, tokenizer = initialize_model()
 
@@ -114,12 +178,16 @@ def main():
     
     # Sidebar Configuration
     st.sidebar.markdown("### Training Parameters")
-    num_generations = st.sidebar.slider("Generations", 1, 20, 5)
-    population_size = st.sidebar.slider("Population Size", 4, 20, 6)
+    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)
+    
+    # Initialize Genetic Algorithm
+    ga = GeneticAlgorithm(model, tokenizer, train_dataset, population_size, mutation_rate)
     
     # Run Training
     if st.button("Start Training"):
-        training_loop(dashboard, model, train_dataset, num_generations, population_size)
+        training_loop(dashboard, ga, num_generations)
 
 if __name__ == "__main__":
     main()