Update app.py

app.py

@@ -2,38 +2,59 @@ import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import DataLoader
-from datasets import load_dataset
+from torch.utils.data import DataLoader, Dataset
 from transformers import AutoTokenizer, get_scheduler
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import io
 from PIL import Image
 import openai
 import time
 
-# ✅ Secure OpenAI API key
+# ✅ Set OpenAI API key from secret
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# ✅ Use GPU if available
+# ✅ Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ✅ Load official PIQA dataset with remote code trust enabled
-dataset = load_dataset("piqa", trust_remote_code=True)
+# ✅ Load PIQA from public GitHub (JSONL)
+dataset = {
+    "train": pd.read_json("https://raw.githubusercontent.com/epfml/Deep_Learning_Projects/master/PIQA/data/train.jsonl", lines=True),
+    "validation": pd.read_json("https://raw.githubusercontent.com/epfml/Deep_Learning_Projects/master/PIQA/data/valid.jsonl", lines=True)
+}
+
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
+# ✅ Tokenization helper
 def tokenize_choices(example):
-    input_0 = tokenizer(example["goal"] + " " + example["sol1"], truncation=True, padding="max_length", max_length=128)
-    input_1 = tokenizer(example["goal"] + " " + example["sol2"], truncation=True, padding="max_length", max_length=128)
+    input_0 = tokenizer(example["goal"] + " " + example["sol1"], truncation=True, padding="max_length", max_length=128, return_tensors="pt")
+    input_1 = tokenizer(example["goal"] + " " + example["sol2"], truncation=True, padding="max_length", max_length=128, return_tensors="pt")
     return {
-        "input_ids_0": input_0["input_ids"],
-        "input_ids_1": input_1["input_ids"],
-        "label": example["label"]
+        "input_ids_0": input_0["input_ids"][0],
+        "input_ids_1": input_1["input_ids"][0],
+        "label": int(example["label"])
     }
 
-dataset = dataset.map(tokenize_choices)
-val_dataset = dataset["validation"].select(range(200)).with_format("torch")
+train_data = [tokenize_choices(row) for _, row in dataset["train"].head(500).iterrows()]
+val_data = [tokenize_choices(row) for _, row in dataset["validation"].head(200).iterrows()]
+
+# ✅ Dataset class
+class PIQADataset(Dataset):
+    def __init__(self, data):
+        self.data = data
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, idx):
+        return {
+            "input_ids_0": self.data[idx]["input_ids_0"],
+            "input_ids_1": self.data[idx]["input_ids_1"],
+            "label": torch.tensor(self.data[idx]["label"])
+        }
+
+train_dataset = PIQADataset(train_data)
+val_dataset = PIQADataset(val_data)
 
 # ✅ EvoTransformer definition
 class EvoTransformer(nn.Module):
@@ -53,7 +74,7 @@ class EvoTransformer(nn.Module):
         x = self.encoder(x)
         return self.classifier(x[:, 0, :]).squeeze(-1)
 
-# ✅ GPT-3.5 response
+# ✅ GPT-3.5 logic
 def gpt35_answer(prompt):
     try:
         response = openai.ChatCompletion.create(
@@ -66,19 +87,17 @@ def gpt35_answer(prompt):
     except Exception as e:
         return f"[Error: {e}]"
 
-# ✅ Training and evaluation function
+# ✅ Main train + compare function
 def train_and_demo(few_shot_size):
     start_time = time.time()
     model = EvoTransformer().to(device)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.AdamW(model.parameters(), lr=5e-5)
 
-    train_set = dataset["train"].select(range(few_shot_size)).with_format("torch")
-    train_loader = DataLoader(train_set, batch_size=8, shuffle=True)
+    loader = DataLoader(train_dataset[:few_shot_size], batch_size=8, shuffle=True)
     val_loader = DataLoader(val_dataset, batch_size=32)
 
-    scheduler = get_scheduler("linear", optimizer=optimizer,
-                              num_warmup_steps=0, num_training_steps=3 * len(train_loader))
+    scheduler = get_scheduler("linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=3 * len(loader))
 
     best_val = 0
     accs = []
@@ -87,7 +106,7 @@ def train_and_demo(few_shot_size):
 
     for epoch in range(3):
         model.train()
-        for batch in train_loader:
+        for batch in loader:
             optimizer.zero_grad()
             x0 = batch["input_ids_0"].to(device)
             x1 = batch["input_ids_1"].to(device)
@@ -122,7 +141,7 @@ def train_and_demo(few_shot_size):
             if early_stop >= patience:
                 break
 
-    # ✅ Accuracy Plot
+    # ✅ Accuracy plot
     fig, ax = plt.subplots()
     ax.plot(accs, marker='o')
     ax.set_title(f"Validation Accuracy ({few_shot_size} examples)")
@@ -133,16 +152,16 @@ def train_and_demo(few_shot_size):
     buf.seek(0)
     img = Image.open(buf)
 
-    # ✅ GPT vs Evo Predictions
+    # ✅ Example comparison with GPT-3.5
     output = ""
     for i in range(2):
-        ex = dataset["validation"][i]
+        ex = dataset["validation"].iloc[i]
         goal = ex["goal"]
         sol1 = ex["sol1"]
         sol2 = ex["sol2"]
 
-        x0 = torch.tensor([ex["input_ids_0"]]).to(device)
-        x1 = torch.tensor([ex["input_ids_1"]]).to(device)
+        x0 = tokenizer(goal + " " + sol1, return_tensors="pt", padding="max_length", max_length=128, truncation=True)["input_ids"].to(device)
+        x1 = tokenizer(goal + " " + sol2, return_tensors="pt", padding="max_length", max_length=128, truncation=True)["input_ids"].to(device)
         l0 = model(x0)
         l1 = model(x1)
         pred_evo = 0 if l0 > l1 else 1
@@ -167,15 +186,15 @@ EvoTransformer v2.1 Configuration:
 
     return img, f"Best Accuracy: {best_val:.4f}", output.strip() + "\n\n" + architecture_info.strip()
 
-# ✅ Gradio interface
+# ✅ Gradio app
 gr.Interface(
     fn=train_and_demo,
-    inputs=gr.Slider(10, 500, step=10, value=50, label="Number of Training Examples"),
+    inputs=gr.Slider(10, 300, step=10, value=50, label="Training Samples"),
     outputs=[
         gr.Image(label="Accuracy Plot"),
         gr.Textbox(label="Best Accuracy"),
         gr.Textbox(label="Evo vs GPT-3.5 Output")
     ],
     title="🧬 EvoTransformer v2.1 Benchmark",
-    description="Train EvoTransformer live on PIQA and compare with GPT-3.5."
+    description="Train EvoTransformer on PIQA and compare predictions against GPT-3.5."
 ).launch()