Update app.py

app.py

@@ -1,200 +1,14 @@
-import os
-import torch
-import torch.nn as nn
-import torch.optim as optim
-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
+from inference import predict
 
-# ✅ Set OpenAI API key from secret
-openai.api_key = os.getenv("OPENAI_API_KEY")
+def compare(goal, sol1, sol2):
+    return predict(goal, sol1, sol2)
 
-# ✅ Device setup
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+demo = gr.Interface(
+    fn=compare,
+    inputs=[gr.Text(label="Goal"), gr.Text(label="Solution 1"), gr.Text(label="Solution 2")],
+    outputs=gr.Text(label="Preferred Solution"),
+    title="EvoTransformer v2.1 - PIQA Commonsense Reasoning"
+)
 
-# ✅ 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, 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"][0],
-        "input_ids_1": input_1["input_ids"][0],
-        "label": int(example["label"])
-    }
-
-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):
-    def __init__(self):
-        super().__init__()
-        self.embedding = nn.Embedding(30522, 384)
-        encoder_layer = nn.TransformerEncoderLayer(d_model=384, nhead=6, dim_feedforward=1024, batch_first=True)
-        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
-        self.classifier = nn.Sequential(
-            nn.Linear(384, 128),
-            nn.ReLU(),
-            nn.Linear(128, 1)
-        )
-
-    def forward(self, input_ids):
-        x = self.embedding(input_ids)
-        x = self.encoder(x)
-        return self.classifier(x[:, 0, :]).squeeze(-1)
-
-# ✅ GPT-3.5 logic
-def gpt35_answer(prompt):
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
-            messages=[{"role": "user", "content": prompt}],
-            max_tokens=20,
-            temperature=0
-        )
-        return response['choices'][0]['message']['content'].strip()
-    except Exception as e:
-        return f"[Error: {e}]"
-
-# ✅ 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)
-
-    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(loader))
-
-    best_val = 0
-    accs = []
-    patience = 2
-    early_stop = 0
-
-    for epoch in range(3):
-        model.train()
-        for batch in loader:
-            optimizer.zero_grad()
-            x0 = batch["input_ids_0"].to(device)
-            x1 = batch["input_ids_1"].to(device)
-            labels = batch["label"].to(device)
-            l0 = model(x0)
-            l1 = model(x1)
-            logits = torch.stack([l0, l1], dim=1)
-            loss = criterion(logits, labels)
-            loss.backward()
-            optimizer.step()
-            scheduler.step()
-
-        model.eval()
-        correct = 0
-        with torch.no_grad():
-            for batch in val_loader:
-                x0 = batch["input_ids_0"].to(device)
-                x1 = batch["input_ids_1"].to(device)
-                labels = batch["label"].to(device)
-                l0 = model(x0)
-                l1 = model(x1)
-                logits = torch.stack([l0, l1], dim=1)
-                preds = torch.argmax(logits, dim=1)
-                correct += (preds == labels).sum().item()
-        acc = correct / len(val_dataset)
-        accs.append(acc)
-        if acc > best_val:
-            best_val = acc
-            early_stop = 0
-        else:
-            early_stop += 1
-            if early_stop >= patience:
-                break
-
-    # ✅ Accuracy plot
-    fig, ax = plt.subplots()
-    ax.plot(accs, marker='o')
-    ax.set_title(f"Validation Accuracy ({few_shot_size} examples)")
-    ax.set_xlabel("Epoch")
-    ax.set_ylabel("Accuracy")
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    img = Image.open(buf)
-
-    # ✅ Example comparison with GPT-3.5
-    output = ""
-    for i in range(2):
-        ex = dataset["validation"].iloc[i]
-        goal = ex["goal"]
-        sol1 = ex["sol1"]
-        sol2 = ex["sol2"]
-
-        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
-        correct_evo = "✅" if pred_evo == ex["label"] else "❌"
-
-        gpt_prompt = f"Q: {goal}\nA) {sol1}\nB) {sol2}\nWhich is more appropriate? Answer with A or B only."
-        gpt_out = gpt35_answer(gpt_prompt)
-        pred_gpt = gpt_out[0].upper()
-        correct_gpt = "✅" if (pred_gpt == 'A' and ex["label"] == 0) or (pred_gpt == 'B' and ex["label"] == 1) else "❌"
-
-        output += f"Q: {goal}\nA) {sol1}\nB) {sol2}\n\nEvoTransformer: {'A' if pred_evo==0 else 'B'} {correct_evo}\nGPT-3.5: {pred_gpt} {correct_gpt}\n\n"
-
-    architecture_info = f"""
-EvoTransformer v2.1 Configuration:
-- Embedding Dim: 384
-- Transformer Layers: 6
-- Attention Heads: 6
-- Feedforward Size: 1024
-- Parameters: ~13M
-- Training Time: {time.time() - start_time:.2f}s
-    """
-
-    return img, f"Best Accuracy: {best_val:.4f}", output.strip() + "\n\n" + architecture_info.strip()
-
-# ✅ Gradio app
-gr.Interface(
-    fn=train_and_demo,
-    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 on PIQA and compare predictions against GPT-3.5."
-).launch()
+demo.launch()