Update app.py

app.py

@@ -1,33 +1,30 @@
 import torch
+import torch.nn as nn
+from config import Net  # Make sure this matches your class name
 import gradio as gr
-import numpy as np
-from config import MLP
 
-# Load model
-model = MLP()
+# Load the model
+model = Net()
 model.load_state_dict(torch.load("pytorch_model.pth", map_location=torch.device("cpu")))
 model.eval()
 
-# Example class names (you can change this)
-class_names = [f"Class {i}" for i in range(8)]
-
-# Prediction function
-def predict(input_vector):
-    input_array = np.array(input_vector).astype(np.float32)
-    if len(input_array) != 1000:
-        return "Error: Input must be 1000 numbers"
-    tensor = torch.tensor(input_array).unsqueeze(0)
+# Define a prediction function
+def predict(inputs):
     with torch.no_grad():
-        output = model(tensor)
-        probs = torch.nn.functional.softmax(output[0], dim=0)
-    return {class_names[i]: float(probs[i]) for i in range(8)}
+        inputs = torch.tensor(inputs).float().unsqueeze(0)  # Add batch dimension
+        output = model(inputs)
+        if isinstance(output, torch.Tensor):
+            return output.squeeze().tolist()
+        return output  # fallback
 
-# Gradio interface
+# Create the Gradio interface
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Textbox(lines=5, placeholder="Enter 1000 comma-separated numbers..."),
-    outputs=gr.Label(num_top_classes=3),
-    title="MLP Vector Classifier"
+    inputs=gr.Textbox(label="Enter comma-separated input values (e.g., 1.2, 3.4, 5.6)"),
+    outputs=gr.Textbox(label="Model Output"),
+    title="PyTorch MLP Classifier"
 )
 
-demo.launch()
+# Launch
+if __name__ == "__main__":
+    demo.launch()