Update app.py

app.py

@@ -16,14 +16,11 @@ import time
 # ✅ Secure OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# ✅ Use CPU or GPU
+# ✅ Use GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ✅ 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"
-})
+# ✅ Load official PIQA dataset with remote code trust enabled
+dataset = load_dataset("piqa", trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
 def tokenize_choices(example):
@@ -69,7 +66,7 @@ def gpt35_answer(prompt):
     except Exception as e:
         return f"[Error: {e}]"
 
-# ✅ Training + Evaluation function
+# ✅ Training and evaluation function
 def train_and_demo(few_shot_size):
     start_time = time.time()
     model = EvoTransformer().to(device)
@@ -125,7 +122,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)")
@@ -136,7 +133,7 @@ def train_and_demo(few_shot_size):
     buf.seek(0)
     img = Image.open(buf)
 
-    # ✅ Show comparison examples
+    # ✅ GPT vs Evo Predictions
     output = ""
     for i in range(2):
         ex = dataset["validation"][i]
@@ -170,7 +167,7 @@ EvoTransformer v2.1 Configuration:
 
     return img, f"Best Accuracy: {best_val:.4f}", output.strip() + "\n\n" + architecture_info.strip()
 
-# ✅ Gradio Interface
+# ✅ Gradio interface
 gr.Interface(
     fn=train_and_demo,
     inputs=gr.Slider(10, 500, step=10, value=50, label="Number of Training Examples"),
@@ -180,5 +177,5 @@ gr.Interface(
         gr.Textbox(label="Evo vs GPT-3.5 Output")
     ],
     title="🧬 EvoTransformer v2.1 Benchmark",
-    description="Train EvoTransformer on PIQA and compare its predictions to GPT-3.5."
+    description="Train EvoTransformer live on PIQA and compare with GPT-3.5."
 ).launch()