Update watchdog.py

watchdog.py

@@ -1,10 +1,10 @@
-from firebase_admin import firestore
 from evo_model import EvoTransformerForClassification
 from transformers import AutoTokenizer
 import torch
-from torch.utils.data import Dataset, DataLoader
 import torch.nn as nn
 import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from firebase_admin import firestore
 
 class EvoDataset(Dataset):
     def __init__(self, texts, labels, tokenizer, max_length=64):
@@ -19,49 +19,53 @@ class EvoDataset(Dataset):
     def __len__(self):
         return len(self.labels)
 
-def train_evo_transformer():
-    db = firestore.client()
-    docs = db.collection("evo_feedback_logs").stream()
+def manual_retrain():
+    try:
+        db = firestore.client()
+        docs = db.collection("evo_feedback_logs").stream()
 
-    goals, solution1, solution2, labels = [], [], [], []
-    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)
+        goals, solution1, solution2, labels = [], [], [], []
+        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)
 
-    if not goals:
-        print("[Retrain Error] No training data found.")
-        return False
+        if not goals:
+            print("[Retrain Error] No 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)
+        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)
 
-    config = {
-        "vocab_size": tokenizer.vocab_size,
-        "d_model": 256,
-        "nhead": 4,
-        "dim_feedforward": 512,
-        "num_hidden_layers": 4
-    }
-    model = EvoTransformerForClassification.from_config_dict(config)
-    model.train()
+        config = {
+            "vocab_size": tokenizer.vocab_size,
+            "d_model": 256,
+            "nhead": 4,
+            "dim_feedforward": 512,
+            "num_hidden_layers": 4
+        }
+        model = EvoTransformerForClassification.from_config_dict(config)
+        model.train()
 
-    optimizer = optim.AdamW(model.parameters(), lr=1e-4)
-    criterion = nn.CrossEntropyLoss()
+        optimizer = optim.AdamW(model.parameters(), lr=1e-4)
+        criterion = nn.CrossEntropyLoss()
 
-    for epoch in range(3):  # quick training
-        for input_ids, label in loader:
-            logits = model(input_ids)
-            loss = criterion(logits, label)
-            loss.backward()
-            optimizer.step()
-            optimizer.zero_grad()
+        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()
 
-    model.save_pretrained("trained_evo")
-    print("✅ Retraining complete.")
-    return True
+        model.save_pretrained("trained_evo")
+        print("✅ Retraining complete.")
+        return True
+    except Exception as e:
+        print(f"[Retrain Error] {e}")
+        return False