Spaces:
Configuration error
Configuration error
File size: 1,485 Bytes
c34c47c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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()
|