Update watchdog.py

watchdog.py

@@ -1,35 +1,67 @@
+from firebase_admin import firestore
+from evo_model import EvoTransformerForClassification
+from transformers import AutoTokenizer
 import torch
-import pandas as pd
-from evo_model import EvoTransformer, train_evo_transformer
-from datasets import load_dataset
-import os
-
-def manual_retrain():
-    try:
-        # Load feedback data from Firestore
-        from google.cloud import firestore
-        db = firestore.Client.from_service_account_json("firebase_key.json")
-        docs = db.collection("evo_feedback_logs").stream()
-        data = [doc.to_dict() for doc in docs if "goal" in doc.to_dict()]
-        if not data:
-            print("No feedback data available.")
-            return False
-
-        # Convert to training format
-        rows = []
-        for d in data:
-            question = d["goal"]
-            option1 = d["sol1"]
-            option2 = d["sol2"]
-            correct = d["correct"]
-            label = 0 if correct == "Solution 1" else 1
-            rows.append((question, option1, option2, label))
-        df = pd.DataFrame(rows, columns=["goal", "sol1", "sol2", "label"])
-
-        # Train the Evo model (minimal epochs to simulate update)
-        train_evo_transformer(df, epochs=1)
-
-        return True
-    except Exception as e:
-        print(f"[Retrain Error] {e}")
+from torch.utils.data import Dataset, DataLoader
+import torch.nn as nn
+import torch.optim as optim
+
+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
+
+    def __getitem__(self, idx):
+        input_ids = torch.tensor(self.encodings["input_ids"][idx])
+        label = torch.tensor(self.labels[idx])
+        return input_ids, label
+
+    def __len__(self):
+        return len(self.labels)
+
+def train_evo_transformer():
+    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)
+
+    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)
+
+    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()
+
+    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()
+
+    model.save_pretrained("trained_evo")
+    print("✅ Retraining complete.")
+    return True