Update backend/merge_engines/evolutionary.py

backend/merge_engines/evolutionary.py

@@ -1,70 +1,85 @@
 import os
-import torch
 import random
-import numpy as np
-from .linear import linear_merge  # reuse for individual creation
-from backend.fitness import evaluate_model
 import tempfile
 import shutil
+import json
+from .linear import merge_models
+from backend.fitness import evaluate_model
 
 def evolutionary_merge(model_a, model_b, dataset_path, output_dir, params):
     pop_size = params.get("population_size", 10)
     mutation_rate = params.get("mutation_rate", 0.1)
-    crossover_method = params.get("crossover", "uniform")  # or "single_point"
+    crossover_method = params.get("crossover", "uniform")  # uniform or arithmetic
     selection = params.get("selection", "tournament")
     generations = params.get("generations", 5)
-    fitness_script = params.get("fitness_script", None)  # custom function string?
-    
+    fitness_higher_better = params.get("fitness_higher_better", False)
+    fitness_script = params.get("fitness_script", None)
+
     # Initialize population: list of alpha values
     population = [random.random() for _ in range(pop_size)]
     best_alpha = None
-    best_score = float("inf") if params.get("fitness_higher_better", False) else -float("inf")
-    
+    best_score = float('-inf') if fitness_higher_better else float('inf')
+
     for gen in range(generations):
         scores = []
         for idx, alpha in enumerate(population):
             temp_dir = tempfile.mkdtemp()
-            linear_merge(model_a, model_b, alpha, temp_dir)
-            score = evaluate_model(temp_dir, dataset_path, fitness_script)
-            scores.append(score)
-            shutil.rmtree(temp_dir)
+            try:
+                # Use merge_models with method='linear' and alpha
+                merge_models(model_a, model_b, temp_dir, method='linear', alpha=alpha)
+                score = evaluate_model(temp_dir, dataset_path, fitness_script)
+                scores.append(score)
+            finally:
+                shutil.rmtree(temp_dir)
+
             # update best
-            if (params.get("fitness_higher_better", False) and score > best_score) or \
-               (not params.get("fitness_higher_better", False) and score < best_score):
-                best_score = score
-                best_alpha = alpha
+            if fitness_higher_better:
+                if score > best_score:
+                    best_score = score
+                    best_alpha = alpha
+            else:
+                if score < best_score:
+                    best_score = score
+                    best_alpha = alpha
+
         # Selection
-        new_pop = []
-        # elitism: keep best
-        new_pop.append(best_alpha)
+        new_pop = [best_alpha]  # elitism: keep best
+
         for _ in range(pop_size - 1):
             if selection == "tournament":
                 i, j = random.sample(range(pop_size), 2)
-                winner = population[i] if scores[i] > scores[j] else population[j] if params.get("fitness_higher_better") else (population[i] if scores[i] < scores[j] else population[j])
+                if fitness_higher_better:
+                    winner = population[i] if scores[i] > scores[j] else population[j]
+                else:
+                    winner = population[i] if scores[i] < scores[j] else population[j]
                 new_pop.append(winner)
             else:
-                # random parent
                 new_pop.append(random.choice(population))
+
         # Crossover & Mutation
         offspring = []
-        for i in range(0, pop_size, 2):
+        for i in range(0, pop_size-1, 2):
             p1 = new_pop[i]
-            p2 = new_pop[(i+1) % pop_size]
+            p2 = new_pop[i+1]
             if crossover_method == "uniform":
                 c1 = p1 if random.random() < 0.5 else p2
                 c2 = p2 if random.random() < 0.5 else p1
             else:  # arithmetic
                 c1 = (p1 + p2) / 2
                 c2 = (p1 + p2) / 2
-            # mutation
-            c1 = min(max(c1 + random.uniform(-mutation_rate, mutation_rate), 0), 1)
-            c2 = min(max(c2 + random.uniform(-mutation_rate, mutation_rate), 0), 1)
+
+            # Mutation
+            c1 = min(max(c1 + random.uniform(-mutation_rate, mutation_rate), 0.0), 1.0)
+            c2 = min(max(c2 + random.uniform(-mutation_rate, mutation_rate), 0.0), 1.0)
             offspring.extend([c1, c2])
+
         population = offspring[:pop_size]
-    
+
     # Final merge with best alpha
-    linear_merge(model_a, model_b, best_alpha, output_dir)
-    # Save fitness info
+    if best_alpha is None:
+        best_alpha = 0.5  # fallback
+    merge_models(model_a, model_b, output_dir, method='linear', alpha=best_alpha)
+
+    # Save info
     with open(os.path.join(output_dir, "merge_info.json"), "w") as f:
-        import json
         json.dump({"best_alpha": best_alpha, "fitness": best_score}, f)