Update watchdog.py

watchdog.py

@@ -1,10 +1,12 @@
+# watchdog.py
+
 import torch
 from transformers import AutoTokenizer
 from evo_model import EvoTransformerForClassification
 from firebase_admin import firestore
 import pandas as pd
 
-# Load tokenizer once
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
 def load_feedback_data():
@@ -24,8 +26,7 @@ def load_feedback_data():
 
 def encode(goal, sol1, sol2):
     prompt = f"Goal: {goal} Option 1: {sol1} Option 2: {sol2}"
-    encoded = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    return encoded.input_ids, encoded.attention_mask
+    return tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
 
 def manual_retrain():
     try:
@@ -36,28 +37,25 @@ def manual_retrain():
 
         model = EvoTransformerForClassification.from_pretrained("trained_model")
         optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-        loss_fn = torch.nn.CrossEntropyLoss()
 
         model.train()
         for _, row in data.sample(frac=1).iterrows():
-            input_ids, attention_mask = encode(row["goal"], row["sol1"], row["sol2"])
-            label = torch.tensor([row["label"]])
-
-            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-            logits = outputs.logits if hasattr(outputs, "logits") else outputs
+            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 logits.ndim == 2 and label.ndim == 1:
-                loss = loss_fn(logits, label)
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
+            if isinstance(outputs, tuple):
+                loss = outputs[0]
             else:
-                print("[Retrain Warning] Shape mismatch, skipping one example.")
+                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