Update retrain_from_feedback

retrain_from_feedback

@@ -1,5 +1,3 @@
-# retrain_from_feedback.py
-
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -7,10 +5,11 @@ from torch.utils.data import DataLoader, Dataset
 from transformers import AutoTokenizer
 
 from evo_architecture import mutate_genome, default_config, log_genome
-from model import EvoTransformerV22  # Must accept dynamic config
+from model import EvoTransformerV22  # Ensure this is compatible with config
 import csv
 import os
 
+# Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class FeedbackDataset(Dataset):
@@ -23,11 +22,13 @@ class FeedbackDataset(Dataset):
         return len(self.samples)
 
     def __getitem__(self, idx):
-        q, o1, o2, ctx, ans, label = self.samples[idx]
-        text = f"{q} [SEP] {o1} [SEP] {o2} [SEP] {ctx}"
-        enc = self.tokenizer(text, padding="max_length", truncation=True, max_length=self.max_len, return_tensors="pt")
+        q, o1, o2, ctx, evo_ans = self.samples[idx]
+        prompt = f"{q} [SEP] {o1} [SEP] {o2} [SEP] {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 = 1 if ans.strip().lower() == o1.strip().lower() else 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)
 
 def load_feedback():
@@ -38,20 +39,19 @@ def load_feedback():
     with open("feedback_log.csv", encoding="utf-8") as f:
         reader = csv.DictReader(f)
         for row in reader:
-            if row["evo_was_correct"].strip().lower() == "yes":
+            if row.get("evo_was_correct", "no").strip().lower() == "yes":
                 data.append([
                     row["question"],
                     row["option1"],
                     row["option2"],
                     row["context"],
-                    row["evo_output"],
-                    "yes"
+                    row["evo_output"].strip()
                 ])
     return data
 
 def build_model(config):
     from model import EvoEncoder
-    class CustomEvo(nn.Module):
+    class EvoClassifier(nn.Module):
         def __init__(self):
             super().__init__()
             self.encoder = EvoEncoder(
@@ -62,14 +62,14 @@ def build_model(config):
                 memory_enabled=config["memory_enabled"]
             )
             self.pool = nn.AdaptiveAvgPool1d(1)
-            self.classifier = nn.Linear(512, 1)
+            self.classifier = nn.Linear(512, 2)  # two-class classification
 
         def forward(self, input_ids):
             x = self.encoder(input_ids)
             x = self.pool(x.transpose(1, 2)).squeeze(-1)
             return self.classifier(x)
 
-    return CustomEvo().to(device)
+    return EvoClassifier().to(device)
 
 def train_evo():
     tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
@@ -86,27 +86,26 @@ def train_evo():
     dataset = FeedbackDataset(tokenizer, data)
     loader = DataLoader(dataset, batch_size=4, shuffle=True)
 
-    loss_fn = nn.BCEWithLogitsLoss()
+    loss_fn = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=1e-4)
 
     for epoch in range(3):
         total_loss, correct = 0, 0
         for input_ids, labels in loader:
-            input_ids, labels = input_ids.to(device), labels.float().to(device)
-            logits = model(input_ids).squeeze(-1)
+            input_ids, labels = input_ids.to(device), labels.to(device)
+            logits = model(input_ids)
             loss = loss_fn(logits, labels)
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
 
             total_loss += loss.item()
-            preds = (torch.sigmoid(logits) > 0.5).long()
-            correct += (preds == labels.long()).sum().item()
+            preds = torch.argmax(logits, dim=1)
+            correct += (preds == labels).sum().item()
 
         acc = correct / len(dataset)
         print(f"✅ Epoch {epoch+1} | Loss={total_loss:.4f} | Acc={acc:.4f}")
 
-    # Save model + genome
     os.makedirs("trained_model", exist_ok=True)
     torch.save(model.state_dict(), "trained_model/evo_retrained.pt")
     log_genome(new_config, acc)