changed the pytorch model

app.py

@@ -1,29 +1,48 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+import torch
 import pandas as pd
-import joblib
 
-# Load the trained model
-model = joblib.load("model.pkl")  # Ensure your model is saved as 'model.pkl'
+# Define the model structure (ensure this matches your model class)
+class AsteroidModel(torch.nn.Module):
+    def __init__(self):
+        super(AsteroidModel, self).__init__()
+        # Define the layers as per your original model architecture
+        self.fc1 = torch.nn.Linear(5, 16)
+        self.fc2 = torch.nn.Linear(16, 8)
+        self.fc3 = torch.nn.Linear(8, 1)
+    
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = torch.sigmoid(self.fc3(x))
+        return x
+
+# Initialize the model and load the saved weights
+model = AsteroidModel()
+model.load_state_dict(torch.load('model.pth'))
+model.eval()  # Set model to evaluation mode
 
 app = FastAPI()
 
-# Add CORS middleware to allow requests from any origin
+# CORS middleware to handle cross-origin requests
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins (adjust if needed)
+    allow_origins=["*"],  # Allow all origins, adjust if needed
     allow_credentials=True,
-    allow_methods=["*"],  # Allow all methods (GET, POST, etc.)
+    allow_methods=["*"],  # Allow all methods
     allow_headers=["*"],  # Allow all headers
 )
 
 @app.post("/predict")
 async def predict(features: dict):
-    # Convert the input into a DataFrame
+    # Convert the input to a tensor
     input_data = pd.DataFrame([features])
+    input_tensor = torch.tensor(input_data.values, dtype=torch.float32)
     
-    # Make prediction using the trained model
-    prediction = model.predict(input_data)
-    
-    return {"is_potentially_hazardous_asteroid": int(prediction[0])}
+    # Make prediction
+    with torch.no_grad():
+        output = model(input_tensor).squeeze()
+        prediction = (output > 0.5).float().item()  # Convert to binary prediction
 
+    return {"is_potentially_hazardous_asteroid": int(prediction)}

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:750fa82060eb6b41c0be0db37365efb3c7ed2af1cdb5095491ec6a9a7ceae8ca
+size 3252

requirements.txt

@@ -1,5 +1,4 @@
 fastapi
 uvicorn
 pandas
-scikit-learn
-joblib
+torch