Update evo_model.py

evo_model.py

@@ -1,9 +1,7 @@
 import torch
 from torch import nn
-from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer
+from transformers import PreTrainedModel, PretrainedConfig
 from model import EvoTransformer  # assumes your core model is in model.py
-from torch.utils.data import DataLoader, Dataset
-import torch.optim as optim
 
 class EvoTransformerConfig(PretrainedConfig):
     model_type = "evo-transformer"
@@ -36,64 +34,21 @@ class EvoTransformerForClassification(PreTrainedModel):
             dim_feedforward=config.dim_feedforward,
             num_layers=config.num_hidden_layers
         )
-        self.classifier = nn.Linear(config.d_model, 2)
+        self.classifier = nn.Linear(config.d_model, 2)  # 2-way classification
 
-    def forward(self, input_ids, attention_mask=None):
-        x = self.model(input_ids)  # (batch_size, seq_len, hidden_size)
-        pooled = x[:, 0, :]  # Take [CLS]-like first token
-        logits = self.classifier(pooled)
+    def forward(self, input_ids):
+        hidden = self.model(input_ids)  # (batch_size, seq_len, d_model)
+        pooled = hidden[:, 0, :]        # Use the first token as a summary
+        logits = self.classifier(pooled)  # (batch_size, 2)
         return logits
 
     def save_pretrained(self, save_directory):
-        torch.save(self.state_dict(), f"{save_directory}/pytorch_model.bin")
+        torch.save(self.model.state_dict(), f"{save_directory}/pytorch_model.bin")
         self.config.save_pretrained(save_directory)
 
     @classmethod
     def from_pretrained(cls, load_directory):
         config = EvoTransformerConfig.from_pretrained(load_directory)
         model = cls(config)
-        model.load_state_dict(torch.load(f"{load_directory}/pytorch_model.bin"))
+        model.model.load_state_dict(torch.load(f"{load_directory}/pytorch_model.bin"))
         return model
-
-# ✅ Add this retraining logic
-def train_evo_transformer(df, epochs=1):
-    class EvoDataset(Dataset):
-        def __init__(self, dataframe, tokenizer):
-            self.df = dataframe
-            self.tokenizer = tokenizer
-
-        def __len__(self):
-            return len(self.df)
-
-        def __getitem__(self, idx):
-            row = self.df.iloc[idx]
-            text = f"{row['goal']} [SEP] {row['sol1']} [SEP] {row['sol2']}"
-            encoding = self.tokenizer(text, truncation=True, padding='max_length', max_length=64, return_tensors='pt')
-            input_ids = encoding['input_ids'].squeeze(0)
-            attention_mask = encoding['attention_mask'].squeeze(0)
-            label = torch.tensor(0 if row['correct'] == 'Solution 1' else 1)
-            return input_ids, attention_mask, label
-
-    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-    config = EvoTransformerConfig()
-    model = EvoTransformerForClassification(config)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    model.train()
-
-    dataset = EvoDataset(df, tokenizer)
-    loader = DataLoader(dataset, batch_size=8, shuffle=True)
-    optimizer = optim.Adam(model.parameters(), lr=2e-5)
-    criterion = nn.CrossEntropyLoss()
-
-    for epoch in range(epochs):
-        for input_ids, attention_mask, labels in loader:
-            input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
-            logits = model(input_ids, attention_mask)
-            loss = criterion(logits, labels)
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-    torch.save(model.state_dict(), "trained_model.pt")
-    return True