Update app.py

app.py

@@ -8,6 +8,8 @@ import pandas as pd
 from utils import compute_features
 from scipy.stats import nbinom
 
+from xgboost import XGBRegressor
+
 
 class NegBinomialModel(nn.Module):
     def __init__(self, in_features):
@@ -22,14 +24,20 @@ class NegBinomialModel(nn.Module):
         return mu.squeeze(), alpha
 
     
-model = NegBinomialModel(16)
-model.load_state_dict(torch.load("model_weights.pt", map_location='cpu'))
+model = NegBinomialModel(12)
+model.load_state_dict(torch.load("model_weights(1).pt", map_location='cpu'))
 model.eval()
 
 # MU_BANKS = 2.6035915713614286
 # STD_BANKS = 3.0158890435512125
 
 
+with open("xgb_model(1).json", "r") as f:
+    params = json.load(f)
+
+xgb_model = XGBRegressor(**params)
+
+
 def predict_score(lat, lon):
     # Convert input to tensor
     # inputs = torch.tensor([[lat, lon]], dtype=torch.float32)
@@ -37,7 +45,7 @@ def predict_score(lat, lon):
     print("[INPUTS]", inputs)
     num_banks = inputs.pop("num_banks_in_radius", 0)
 
-    inputs = torch.tensor([lat,lon] + list(inputs.values()), dtype=torch.float32)
+    inputs = torch.tensor(list(inputs.values()), dtype=torch.float32)
 
     # Get model output
     with torch.no_grad():
@@ -46,6 +54,8 @@ def predict_score(lat, lon):
     # Unpack into respective values
     mu_pred = mu_pred.numpy().flatten()
 
+    mu_pred2 = xgb_model.predict(inputs.numpy())
+
     # r = 1/alpha
     # p = r / (r + mu_pred)
 
@@ -76,6 +86,7 @@ def predict_score(lat, lon):
         round(float(score), 3),
         num_banks,
         round(float(mu_pred), 3),
+        round(float(mu_pred2), 3),
         # round(float(log_score),3)
         # "Normal Score": round(float(normal_score), 3),
     )
@@ -90,7 +101,8 @@ interface = gr.Interface(
     outputs=[
         gr.Number(label="Score (0 - 100)"),
         gr.Number(label="Number of Current Banks"),
-        gr.Number(label="Number of Ideal Banks"),
+        gr.Number(label="Number of Ideal Banks (Negative Binomial)"),
+        gr.Number(label="Number of Ideal Banks (XGBoost)"),
         # gr.Number(label="Log Score Probability"),
     ],
     title="Bank Location Scoring Model",