Update watchdog.py

watchdog.py

@@ -1,61 +1,56 @@
-# watchdog.py
-
+import os
 import torch
-from transformers import AutoTokenizer
-from evo_model import EvoTransformerForClassification
-from firebase_admin import firestore
-import pandas as pd
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-
-def load_feedback_data():
-    db = firestore.client()
-    docs = db.collection("evo_feedback_logs").stream()
-    data = []
+import firebase_admin
+from firebase_admin import credentials, firestore
+from model import SimpleEvoModel
+
+# Initialize Firebase if not already initialized
+if not firebase_admin._apps:
+    cred = credentials.Certificate("firebase_key.json")
+    firebase_admin.initialize_app(cred)
+
+db = firestore.client()
+
+def fetch_training_data():
+    logs_ref = db.collection("evo_feedback")
+    docs = logs_ref.stream()
+    
+    inputs, labels = [], []
     for doc in docs:
-        d = doc.to_dict()
-        if all(k in d for k in ["goal", "solution_1", "solution_2", "correct_answer"]):
-            data.append((
-                d["goal"],
-                d["solution_1"],
-                d["solution_2"],
-                0 if d["correct_answer"] == "Solution 1" else 1
-            ))
-    return pd.DataFrame(data, columns=["goal", "sol1", "sol2", "label"])
-
-def encode(goal, sol1, sol2):
-    prompt = f"Goal: {goal} Option 1: {sol1} Option 2: {sol2}"
-    return tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-
-def manual_retrain():
-    try:
-        data = load_feedback_data()
-        if data.empty:
-            print("[Retrain Error] No training data found.")
-            return False
-
-        model = EvoTransformerForClassification.from_pretrained("trained_model")
-        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-
-        model.train()
-        for _, row in data.sample(frac=1).iterrows():
-            encoded = encode(row["goal"], row["sol1"], row["sol2"])
-            labels = torch.tensor([row["label"]])
-            outputs = model(input_ids=encoded["input_ids"], attention_mask=encoded["attention_mask"], labels=labels)
-
-            if isinstance(outputs, tuple):
-                loss = outputs[0]
-            else:
-                loss = outputs
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-        model.save_pretrained("trained_model")
-        print("✅ Evo retrained and saved.")
-        return True
-    except Exception as e:
-        print(f"[Retrain Error] {e}")
-        return False
+        data = doc.to_dict()
+        goal = 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
+            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)
+
+def retrain_and_save():
+    X, y = fetch_training_data()
+    if len(X) < 2:
+        print("⚠️ Not enough training data.")
+        return
+
+    model = SimpleEvoModel()
+    loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    for epoch in range(5):
+        optimizer.zero_grad()
+        output = model(X)
+        loss = loss_fn(output, y)
+        loss.backward()
+        optimizer.step()
+    
+    # Save retrained model to trained_model/
+    os.makedirs("trained_model", exist_ok=True)
+    torch.save(model.state_dict(), "trained_model/pytorch_model.bin")
+    print("✅ EvoTransformer retrained and saved to trained_model/")
+
+if __name__ == "__main__":
+    retrain_and_save()