Update watchdog.py

watchdog.py

@@ -5,6 +5,9 @@ from firebase_admin import credentials, firestore
 from evo_model import EvoTransformerForClassification, EvoTransformerConfig
 from transformers import BertTokenizer
 
+from init_model import load_model
+from dashboard import evolution_accuracy_plot
+
 # Initialize Firebase if not already initialized
 if not firebase_admin._apps:
     cred = credentials.Certificate("firebase_key.json")
@@ -44,35 +47,54 @@ def fetch_training_data(tokenizer):
         torch.tensor(labels, dtype=torch.long)
     )
 
-def retrain_and_save():
-    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-    input_ids, attention_masks, labels = fetch_training_data(tokenizer)
+def get_architecture_summary(model):
+    summary = {
+        "Layers": getattr(model, "num_layers", "N/A"),
+        "Attention Heads": getattr(model, "num_heads", "N/A"),
+        "FFN Dim": getattr(model, "ffn_dim", "N/A"),
+        "Memory Enabled": getattr(model, "use_memory", "N/A"),
+    }
+    return "\n".join(f"{k}: {v}" for k, v in summary.items())
 
-    if input_ids is None or len(input_ids) < 2:
-        print("⚠️ Not enough training data.")
-        return
+def retrain_model():
+    try:
+        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        input_ids, attention_masks, labels = fetch_training_data(tokenizer)
 
-    config = EvoTransformerConfig()
-    model = EvoTransformerForClassification(config)
-    model.train()
+        if input_ids is None or len(input_ids) < 2:
+            return "⚠️ Not enough data to retrain.", None, "Please log more feedback first."
 
-    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
-    loss_fn = torch.nn.CrossEntropyLoss()
+        config = EvoTransformerConfig()
+        model = EvoTransformerForClassification(config)
+        model.train()
 
-    for epoch in range(3):
-        optimizer.zero_grad()
-        outputs = model(input_ids, attention_mask=attention_masks)
-        loss = loss_fn(outputs, labels)
-        loss.backward()
-        optimizer.step()
-        print(f"Epoch {epoch+1}: Loss = {loss.item():.4f}")
+        optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
+        loss_fn = torch.nn.CrossEntropyLoss()
 
-    os.makedirs("trained_model", exist_ok=True)
-    model.save_pretrained("trained_model")
-    print("✅ EvoTransformer retrained and saved to trained_model/")
+        for epoch in range(3):
+            optimizer.zero_grad()
+            outputs = model(input_ids, attention_mask=attention_masks)
+            loss = loss_fn(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            print(f"Epoch {epoch+1}: Loss = {loss.item():.4f}")
 
-if __name__ == "__main__":
-    retrain_and_save()
+        os.makedirs("trained_model", exist_ok=True)
+        model.save_pretrained("trained_model")
+
+        print("✅ EvoTransformer retrained and saved.")
+
+        # Reload the updated model for summary + plot
+        updated_model = load_model()
+        arch_text = get_architecture_summary(updated_model)
+        plot = evolution_accuracy_plot()
 
-# Alias to match expected import
-retrain_model = retrain_and_save
+        return arch_text, plot, "✅ EvoTransformer retrained successfully!"
+
+    except Exception as e:
+        print(f"❌ Retraining failed: {e}")
+        return "❌ Error", None, f"Retrain failed: {e}"
+
+# Support direct script run
+if __name__ == "__main__":
+    retrain_model()