Update evo_transformer.py

evo_transformer.py

@@ -1,9 +1,8 @@
-# evo_transformer.py
 import random
+import matplotlib.pyplot as plt
 
 class EvoTransformer:
     def __init__(self, config=None):
-        # Initialize with default or passed config
         self.config = config or {
             "layers": 4,
             "attention_heads": 4,
@@ -16,7 +15,7 @@ class EvoTransformer:
     def mutate(self):
         new_config = self.config.copy()
         trait = random.choice(list(new_config.keys()))
-        
+
         if trait == "layers":
             new_config[trait] = max(1, new_config[trait] + random.choice([-1, 1]))
         elif trait == "attention_heads":
@@ -31,7 +30,7 @@ class EvoTransformer:
         self.config = new_config
         self.history.append(new_config.copy())
 
-    def evolve(self, generations=3):
+    def evolve(self, generations=5):
         for _ in range(generations):
             self.mutate()
 
@@ -39,10 +38,38 @@ class EvoTransformer:
         return self.history
 
     def evaluate(self):
-        # Simulate an accuracy score for demo purposes
+        # Simulated accuracy for demo
         score = round(random.uniform(0.85, 0.95), 4)
-        return {"accuracy": score, "params": self.estimate_params()}
+        return {
+            "accuracy": score,
+            "params": self.estimate_params()
+        }
 
     def estimate_params(self):
-        # Simulated parameter count based on config
-        return 10 + self.config["layers"] * self.config["ffn_dim"] * 0.001
+        return round(10 + self.config["layers"] * self.config["ffn_dim"] * 0.001, 2)
+
+    def plot_evolution(self):
+        layers = [cfg["layers"] for cfg in self.history]
+        heads = [cfg["attention_heads"] for cfg in self.history]
+        ffn_dims = [cfg["ffn_dim"] for cfg in self.history]
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(layers, label="Layers", marker='o', color='orange')
+        plt.plot(heads, label="Attention Heads", marker='s', color='blue')
+        plt.plot(ffn_dims, label="FFN Dim", marker='^', color='green')
+        plt.xlabel("Generation")
+        plt.ylabel("Value")
+        plt.title("Evolution of Traits")
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+        plt.show()
+
+
+# Run test locally
+if __name__ == "__main__":
+    evo = EvoTransformer()
+    evo.evolve(generations=8)
+    evo.plot_evolution()
+    print("Final Config:", evo.config)
+    print("Evaluation:", evo.evaluate())