EvoTransformer-Demo / evo_transformer.py
HemanM's picture
Update evo_transformer.py
304a057 verified
raw
history blame
2.24 kB
import random
import json
import csv
import io
class EvoTransformer:
def __init__(self):
self.history = []
self.current = {}
def reset(self):
self.history = []
self.current = {}
def mutate(self, config):
config["layers"] = max(1, config["layers"] + random.choice([-1, 1]))
config["attention_heads"] = max(1, config["attention_heads"] + random.choice([-1, 1]))
config["ffn_dim"] = max(64, config["ffn_dim"] + random.choice([-256, 256]))
config["dropout"] = min(0.5, max(0.01, config["dropout"] + random.choice([-0.02, 0.02])))
config["memory"] = random.choice([True, False])
return config
def simulate_accuracy(self, config):
return round(random.uniform(70, 90), 2)
def estimate_parameters(self, config):
return round(
config["layers"] * config["attention_heads"] * config["ffn_dim"] / 1e6,
2
)
def evolve(self, generations):
base_config = {
"layers": 4,
"attention_heads": 4,
"ffn_dim": 1024,
"dropout": 0.1,
"memory": False
}
for _ in range(generations):
base_config = self.mutate(base_config.copy())
acc = self.simulate_accuracy(base_config)
params = self.estimate_parameters(base_config)
record = {"traits": base_config.copy(), "accuracy": acc, "parameters": params}
self.history.append(record)
self.current = self.history[-1]
def get_best_config(self):
return self.current
def export_csv(self):
output = io.StringIO()
writer = csv.writer(output)
writer.writerow(["Generation", "Layers", "Attention Heads", "FFN Dim", "Dropout", "Memory", "Accuracy", "Params"])
for i, entry in enumerate(self.history):
t = entry["traits"]
writer.writerow([
i+1,
t["layers"], t["attention_heads"], t["ffn_dim"],
t["dropout"], t["memory"],
entry["accuracy"], entry["parameters"]
])
return output.getvalue()
def export_json(self):
return json.dumps(self.history, indent=2)