import torch
import streamlit as st
import numpy as np

class Net(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(input_size, hidden_size)
        self.relu = torch.nn.ReLU()
        self.output = torch.nn.Linear(hidden_size, output_size)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        hidden = self.hidden(x)
        relu = self.relu(hidden)
        output = self.output(relu)
        output = self.sigmoid(output)
        return output

def load_model(path):
    model = Net(2, 5, 1)
    model.load_state_dict(torch.load(path))
    return model

def predict(model, input_data):
    with torch.no_grad():
        output = model(input_data)
    return output.numpy()

def main():
    st.title("PyTorch Model Predictor")

    uploaded_file = st.file_uploader("Choose a PyTorch model file (.pt)", type="pt")

    if uploaded_file is not None:
        model = load_model(uploaded_file)
        st.success("Model loaded successfully.")

        st.header("Make a Prediction")
        input_data = np.array([st.number_input("Input 1"), st.number_input("Input 2")])
        if st.button("Predict"):
            prediction = predict(model, torch.from_numpy(input_data).float().to('cpu'))
            st.write("Prediction:", prediction.item())
    else:
        st.warning("Please upload a PyTorch model file (.pt).")

if __name__ == "__main__":
    main()