Update retrain_from_feedback.py

retrain_from_feedback.py

@@ -3,15 +3,14 @@ import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, Dataset
 from transformers import AutoTokenizer
-
 from evo_architecture import mutate_genome, default_config, log_genome
 from evo_model import EvoTransformerV22
-import csv
-import os
+import csv, os
 
-# Device setup
+# 💻 Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# 📦 Dataset built from feedback CSV
 class FeedbackDataset(Dataset):
     def __init__(self, tokenizer, data, max_len=128):
         self.tokenizer = tokenizer
@@ -22,15 +21,14 @@ class FeedbackDataset(Dataset):
         return len(self.samples)
 
     def __getitem__(self, idx):
-        q, o1, o2, ctx, evo_ans = self.samples[idx]
-        prompt = f"{q} [SEP] {o1} [SEP] {o2} [SEP] {ctx}"
+        item = self.samples[idx]
+        q, o1, o2, ctx, label = item
+        prompt = f"{q} [SEP] {o1} [SEP] {o2} [CTX] {ctx}"
         enc = self.tokenizer(prompt, padding="max_length", truncation=True, max_length=self.max_len, return_tensors="pt")
         input_ids = enc["input_ids"].squeeze(0)
-
-        # Label: 0 if Evo picked option1, else 1
-        label = 0 if evo_ans.strip().lower() == o1.strip().lower() else 1
         return input_ids, torch.tensor(label)
 
+# 🧠 Load feedback data
 def load_feedback():
     data = []
     if not os.path.exists("feedback_log.csv"):
@@ -39,16 +37,30 @@ def load_feedback():
     with open("feedback_log.csv", encoding="utf-8") as f:
         reader = csv.DictReader(f)
         for row in reader:
-            if row.get("evo_was_correct", "no").strip().lower() == "yes":
-                data.append([
-                    row["question"],
-                    row["option1"],
-                    row["option2"],
-                    row["context"],
-                    row["evo_output"].strip()
-                ])
+            q = row["question"]
+            o1 = row["option1"]
+            o2 = row["option2"]
+            ctx = row["context"]
+            evo_out = row["evo_output"].strip()
+            vote = row.get("user_preference", "").lower()
+            evo_correct = row.get("evo_was_correct", "").lower()
+
+            # Priority 1: user vote
+            if vote == "evo":
+                label = 1
+            elif vote == "gpt":
+                label = 0
+            # Priority 2: evo correctness
+            elif evo_correct == "yes":
+                label = 1
+            else:
+                continue  # skip uncertain rows
+
+            # Label 1 means Evo was correct/preferred
+            data.append([q, o1, o2, ctx, label])
     return data
 
+# 🔧 Evo model builder from config
 def build_model(config):
     from model import EvoEncoder
     class EvoClassifier(nn.Module):
@@ -62,7 +74,7 @@ def build_model(config):
                 memory_enabled=config["memory_enabled"]
             )
             self.pool = nn.AdaptiveAvgPool1d(1)
-            self.classifier = nn.Linear(512, 2)  # two-class classification
+            self.classifier = nn.Linear(512, 2)
 
         def forward(self, input_ids):
             x = self.encoder(input_ids)
@@ -71,6 +83,7 @@ def build_model(config):
 
     return EvoClassifier().to(device)
 
+# 🔁 Train Evo on feedback
 def train_evo():
     tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
     data = load_feedback()
@@ -109,7 +122,8 @@ def train_evo():
     os.makedirs("trained_model", exist_ok=True)
     torch.save(model.state_dict(), "trained_model/evo_retrained.pt")
     log_genome(new_config, acc)
-    print("✅ Model saved and genome logged.")
+    print("✅ Evo retrained and genome logged.")
 
+# 🔁 Entry point
 if __name__ == "__main__":
     train_evo()