Update watchdog.py

watchdog.py

@@ -1,47 +1,44 @@
-from evo_model import EvoTransformerForClassification
-from transformers import AutoTokenizer
+# watchdog.py
+
 import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import Dataset, DataLoader
-from firebase_admin import firestore
+from evo_model import EvoTransformerForClassification, EvoTransformerConfig
+from transformers import BertTokenizer
+import firebase_admin
+from firebase_admin import credentials, firestore
+import os
+from datetime import datetime
 
-class EvoDataset(Dataset):
-    def __init__(self, texts, labels, tokenizer, max_length=64):
-        self.encodings = tokenizer(texts, truncation=True, padding=True, max_length=max_length)
-        self.labels = labels
+# ✅ Load tokenizer
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 
-    def __getitem__(self, idx):
-        input_ids = torch.tensor(self.encodings["input_ids"][idx])
-        label = torch.tensor(self.labels[idx])
-        return input_ids, label
+# ✅ Init Firebase
+if not firebase_admin._apps:
+    cred = credentials.Certificate("firebase_key.json")
+    firebase_admin.initialize_app(cred)
 
-    def __len__(self):
-        return len(self.labels)
+db = firestore.client()
 
 def manual_retrain():
     try:
-        db = firestore.client()
+        # 🔁 Fetch feedback logs
         docs = db.collection("evo_feedback_logs").stream()
-
-        goals, solution1, solution2, labels = [], [], [], []
+        data = []
         for doc in docs:
             d = doc.to_dict()
             if all(k in d for k in ["goal", "solution_1", "solution_2", "correct_answer"]):
-                goals.append(d["goal"])
-                solution1.append(d["solution_1"])
-                solution2.append(d["solution_2"])
-                labels.append(0 if d["correct_answer"] == "Solution 1" else 1)
+                label = 0 if d["correct_answer"] == "Solution 1" else 1
+                combined = f"{d['goal']} [SEP] {d['solution_1']} [SEP] {d['solution_2']}"
+                data.append((combined, label))
 
-        if not goals:
-            print("[Retrain Error] No training data found.")
+        if not data:
+            print("❌ No valid training data found.")
             return False
 
-        tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-        texts = [f"{g} [SEP] {s1} [SEP] {s2}" for g, s1, s2 in zip(goals, solution1, solution2)]
-        dataset = EvoDataset(texts, labels, tokenizer)
-        loader = DataLoader(dataset, batch_size=4, shuffle=True)
+        # ✅ Tokenize
+        inputs = tokenizer([x[0] for x in data], padding=True, truncation=True, return_tensors="pt")
+        labels = torch.tensor([x[1] for x in data])
 
+        # ✅ Load config + model
         config = {
             "vocab_size": tokenizer.vocab_size,
             "d_model": 256,
@@ -49,23 +46,28 @@ def manual_retrain():
             "dim_feedforward": 512,
             "num_hidden_layers": 4
         }
-        model = EvoTransformerForClassification.from_config_dict(config)
-        model.train()
+        model_config = EvoTransformerConfig(**config)
+        model = EvoTransformerForClassification(model_config)
 
-        optimizer = optim.AdamW(model.parameters(), lr=1e-4)
-        criterion = nn.CrossEntropyLoss()
+        # ✅ Loss + optimizer
+        criterion = torch.nn.CrossEntropyLoss()
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
 
+        # ✅ Train (simple 3-epoch fine-tune)
+        model.train()
         for epoch in range(3):
-            for input_ids, label in loader:
-                logits = model(input_ids)
-                loss = criterion(logits, label)
-                loss.backward()
-                optimizer.step()
-                optimizer.zero_grad()
+            optimizer.zero_grad()
+            outputs = model(inputs["input_ids"])
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            print(f"[Epoch {epoch+1}] Loss: {loss.item():.4f}")
 
-        model.save_pretrained("trained_evo")
-        print("✅ Retraining complete.")
+        # ✅ Save model
+        torch.save(model.state_dict(), "trained_model.pt")
+        print("✅ Evo updated via retrain from feedback!")
         return True
+
     except Exception as e:
         print(f"[Retrain Error] {e}")
         return False