Update init_model.py

init_model.py

@@ -1,14 +1,39 @@
-import os
 import torch
-from evo_model import EvoTransformerConfig, EvoTransformerForClassification
+from evo_model import EvoTransformerForClassification, EvoTransformerConfig
+from transformers import AutoTokenizer
+from torch.utils.data import DataLoader, TensorDataset
+import os
+
+def retrain_model():
+    print("🔄 Starting Evo retrain...")
+
+    # Sample retraining data
+    examples = [
+        "Goal: House on fire. Option 1: Exit house. Option 2: Stay in house.",
+        "Goal: Wet floor. Option 1: Walk slowly. Option 2: Run fast.",
+        "Goal: Loud music. Option 1: Turn it down. Option 2: Ignore it."
+    ]
+    labels = [0, 0, 0]  # Option 1 is correct
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+    model = EvoTransformerForClassification(EvoTransformerConfig())
+
+    inputs = tokenizer(examples, padding=True, truncation=True, return_tensors="pt")
+    labels_tensor = torch.tensor(labels)
+
+    dataset = TensorDataset(inputs["input_ids"], inputs["attention_mask"], labels_tensor)
+    dataloader = DataLoader(dataset, batch_size=2)
 
-# Ensure folder
-os.makedirs("trained_model", exist_ok=True)
+    model.train()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
 
-# Create config and model
-config = EvoTransformerConfig()
-model = EvoTransformerForClassification(config)
+    for epoch in range(2):
+        for input_ids, attention_mask, labels_batch in dataloader:
+            optimizer.zero_grad()
+            loss, _ = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels_batch)
+            loss.backward()
+            optimizer.step()
 
-# Save config and model in Hugging Face format
-model.save_pretrained("trained_model")
-print("✅ EvoTransformer initial model + config saved to 'trained_model/'")
+    os.makedirs("trained_model", exist_ok=True)
+    model.save_pretrained("trained_model")
+    print("✅ Evo retrained and saved.")