Update watchdog.py

watchdog.py

@@ -2,7 +2,8 @@ import os
 import torch
 import firebase_admin
 from firebase_admin import credentials, firestore
-from model import SimpleEvoModel
+from evo_model import EvoTransformerForClassification, EvoTransformerConfig
+from transformers import BertTokenizer
 
 # Initialize Firebase if not already initialized
 if not firebase_admin._apps:
@@ -11,45 +12,63 @@ if not firebase_admin._apps:
 
 db = firestore.client()
 
-def fetch_training_data():
+def fetch_training_data(tokenizer):
     logs_ref = db.collection("evo_feedback")
     docs = logs_ref.stream()
-    
-    inputs, labels = [], []
+
+    input_ids, attention_masks, labels = [], [], []
     for doc in docs:
         data = doc.to_dict()
-        goal = data.get("prompt", "")
+        prompt = data.get("prompt", "")
         winner = data.get("winner", "")
-        if winner:
-            # Simulated encoding
-            vector = [float(ord(c) % 256) / 255.0 for c in (goal + winner)]
-            vector = vector[:768] + [0.0] * max(0, 768 - len(vector))  # pad/truncate
+        if winner and prompt:
+            text = prompt + " [SEP] " + winner
+            encoding = tokenizer(
+                text,
+                truncation=True,
+                padding="max_length",
+                max_length=128,
+                return_tensors="pt"
+            )
+            input_ids.append(encoding["input_ids"][0])
+            attention_masks.append(encoding["attention_mask"][0])
             label = 0 if "1" in winner else 1
-            inputs.append(vector)
             labels.append(label)
-    
-    return torch.tensor(inputs, dtype=torch.float32), torch.tensor(labels, dtype=torch.long)
+
+    if not input_ids:
+        return None, None, None
+
+    return (
+        torch.stack(input_ids),
+        torch.stack(attention_masks),
+        torch.tensor(labels, dtype=torch.long)
+    )
 
 def retrain_and_save():
-    X, y = fetch_training_data()
-    if len(X) < 2:
+    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+    input_ids, attention_masks, labels = fetch_training_data(tokenizer)
+
+    if input_ids is None or len(input_ids) < 2:
         print("⚠️ Not enough training data.")
         return
 
-    model = SimpleEvoModel()
+    config = EvoTransformerConfig()
+    model = EvoTransformerForClassification(config)
+    model.train()
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
     loss_fn = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
-    for epoch in range(5):
+    for epoch in range(3):
         optimizer.zero_grad()
-        output = model(X)
-        loss = loss_fn(output, y)
+        outputs = model(input_ids, attention_mask=attention_masks)
+        loss = loss_fn(outputs, labels)
         loss.backward()
         optimizer.step()
-    
-    # Save retrained model to trained_model/
+        print(f"Epoch {epoch+1}: Loss = {loss.item():.4f}")
+
     os.makedirs("trained_model", exist_ok=True)
-    torch.save(model.state_dict(), "trained_model/pytorch_model.bin")
+    model.save_pretrained("trained_model")
     print("✅ EvoTransformer retrained and saved to trained_model/")
 
 if __name__ == "__main__":