Update watchdog.py

watchdog.py

@@ -4,7 +4,7 @@ from evo_model import EvoTransformerForClassification
 from firebase_admin import firestore
 import pandas as pd
 
-# Load tokenizer
+# Load tokenizer once
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
 def load_feedback_data():
@@ -24,7 +24,8 @@ def load_feedback_data():
 
 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).input_ids
+    encoded = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    return encoded.input_ids, encoded.attention_mask
 
 def manual_retrain():
     try:
@@ -39,15 +40,12 @@ def manual_retrain():
 
         model.train()
         for _, row in data.sample(frac=1).iterrows():
-            inputs = encode(row["goal"], row["sol1"], row["sol2"])
+            input_ids, attention_mask = encode(row["goal"], row["sol1"], row["sol2"])
             label = torch.tensor([row["label"]])
-            outputs = model(inputs)
-            if isinstance(outputs, tuple):
-                logits = outputs[0]
-            elif hasattr(outputs, "logits"):
-                logits = outputs.logits
-            else:
-                logits = outputs
+
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            logits = outputs.logits if hasattr(outputs, "logits") else outputs
+
             if logits.ndim == 2 and label.ndim == 1:
                 loss = loss_fn(logits, label)
                 optimizer.zero_grad()
@@ -59,6 +57,7 @@ def manual_retrain():
         model.save_pretrained("trained_model")
         print("✅ Evo retrained and saved.")
         return True
+
     except Exception as e:
         print(f"[Retrain Error] {e}")
         return False