Update app.py

app.py

@@ -13,14 +13,17 @@ from PIL import Image
 import openai
 import time
 
-# ✅ Secure API key
+# ✅ Secure OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# ✅ Device config
+# ✅ Use CPU or GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ✅ Load PIQA dataset
-dataset = load_dataset("piqa")
+# ✅ Load PIQA from Hugging Face JSON (safe for Spaces)
+dataset = load_dataset("json", data_files={
+    "train": "https://huggingface.co/datasets/AI-Sweden/piqa-downsampled/resolve/main/train.json",
+    "validation": "https://huggingface.co/datasets/AI-Sweden/piqa-downsampled/resolve/main/validation.json"
+})
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
 def tokenize_choices(example):
@@ -35,7 +38,7 @@ def tokenize_choices(example):
 dataset = dataset.map(tokenize_choices)
 val_dataset = dataset["validation"].select(range(200)).with_format("torch")
 
-# ✅ EvoTransformer architecture
+# ✅ EvoTransformer definition
 class EvoTransformer(nn.Module):
     def __init__(self):
         super().__init__()
@@ -53,7 +56,7 @@ class EvoTransformer(nn.Module):
         x = self.encoder(x)
         return self.classifier(x[:, 0, :]).squeeze(-1)
 
-# ✅ GPT-3.5 call
+# ✅ GPT-3.5 response
 def gpt35_answer(prompt):
     try:
         response = openai.ChatCompletion.create(
@@ -66,7 +69,7 @@ def gpt35_answer(prompt):
     except Exception as e:
         return f"[Error: {e}]"
 
-# ✅ Evo vs GPT logic
+# ✅ Training + Evaluation function
 def train_and_demo(few_shot_size):
     start_time = time.time()
     model = EvoTransformer().to(device)
@@ -133,7 +136,7 @@ def train_and_demo(few_shot_size):
     buf.seek(0)
     img = Image.open(buf)
 
-    # ✅ Sample output vs GPT-3.5
+    # ✅ Show comparison examples
     output = ""
     for i in range(2):
         ex = dataset["validation"][i]
@@ -167,15 +170,15 @@ EvoTransformer v2.1 Configuration:
 
     return img, f"Best Accuracy: {best_val:.4f}", output.strip() + "\n\n" + architecture_info.strip()
 
-# ✅ Gradio app
+# ✅ Gradio Interface
 gr.Interface(
     fn=train_and_demo,
-    inputs=gr.Slider(10, 500, step=10, value=50, label="Training Samples"),
+    inputs=gr.Slider(10, 500, step=10, value=50, label="Number of Training Examples"),
     outputs=[
-        gr.Image(label="Validation Accuracy Plot"),
+        gr.Image(label="Accuracy Plot"),
         gr.Textbox(label="Best Accuracy"),
-        gr.Textbox(label="Evo vs GPT-3.5: Predictions & Architecture")
+        gr.Textbox(label="Evo vs GPT-3.5 Output")
     ],
-    title="🧬 EvoTransformer v2.1 vs GPT-3.5",
-    description="Train EvoTransformer on PIQA and benchmark it live against GPT-3.5."
+    title="🧬 EvoTransformer v2.1 Benchmark",
+    description="Train EvoTransformer on PIQA and compare its predictions to GPT-3.5."
 ).launch()