import torch
import torch.nn as nn
from config import MLP
import gradio as gr

# Load the mode
model = MLP()
model.load_state_dict(torch.load("pytorch_model.pth", map_location=torch.device("cpu")))
model.eval()

# Define a prediction function
def predict(inputs):
    with torch.no_grad():
        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

# Create the Gradio interface
demo = gr.Interface(
    fn=predict,
    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"
)

# Launch
if __name__ == "__main__":
    demo.launch()