Update app.py

app.py

@@ -11,32 +11,32 @@ from scipy.stats import nbinom
 from xgboost import XGBRegressor
 import json
 
-# class NegBinomialModel(nn.Module):
-#     def __init__(self, in_features):
-#         super().__init__()
-#         self.linear = nn.Linear(in_features, 1)
-#         self.alpha = nn.Parameter(torch.tensor(0.5))
-
-#     def forward(self, x):
-#         # safer activation than exp()
-#         mu = torch.exp(torch.clamp(self.linear(x), min=-5, max=5))
-#         alpha = torch.clamp(self.alpha, min=1e-3, max=10)
-#         return mu.squeeze(), alpha
+class NegBinomialModel(nn.Module):
+    def __init__(self, in_features):
+        super().__init__()
+        self.linear = nn.Linear(in_features, 1)
+        self.alpha = nn.Parameter(torch.tensor(0.5))
+
+    def forward(self, x):
+        # safer activation than exp()
+        mu = torch.exp(torch.clamp(self.linear(x), min=-5, max=5))
+        alpha = torch.clamp(self.alpha, min=1e-3, max=10)
+        return mu.squeeze(), alpha
 
     
-# model = NegBinomialModel(12)
-# model.load_state_dict(torch.load("model_weights(1).pt", map_location='cpu'))
-# model.eval()
+model = NegBinomialModel(10)
+model.load_state_dict(torch.load("model_weights_2.pt", map_location='cpu'))
+model.eval()
 
-# MU_BANKS = 2.6035915713614286
-# STD_BANKS = 3.0158890435512125
+MU_BANKS = 2.6035915713614286
+STD_BANKS = 3.0158890435512125
 
 
 # with open("xgb_model(1).json", "r") as f:
 #     params = json.load(f)
 
 xgb_model = XGBRegressor()
-xgb_model.load_model("xgb_models.json")
+xgb_model.load_model("xgb_model_2.json")
 
 def predict_score(lat, lon, api_key):
     # Convert input to tensor
@@ -49,10 +49,10 @@ def predict_score(lat, lon, api_key):
 
     inputs = torch.tensor(list(inputs.values()), dtype=torch.float32)
 
-    # with torch.no_grad():
-    #     mu_pred, alpha = model(inputs)
+    with torch.no_grad():
+        mu_pred, alpha = model(inputs)
 
-    # mu_pred = mu_pred.numpy().flatten()
+    mu_pred = mu_pred.numpy().flatten()
 
     mu_pred2 = xgb_model.predict(inputs.unsqueeze(0).numpy())
 
@@ -89,7 +89,7 @@ def predict_score(lat, lon, api_key):
     return (
         round(float(score), 3),
         num_banks,
-        # round(float(mu_pred), 3),
+        round(float(mu_pred), 3),
         round(float(mu_pred2), 3),
         # round(float(log_score),3)
         # "Normal Score": round(float(normal_score), 3),
@@ -110,7 +110,7 @@ interface = gr.Interface(
     outputs=[
         gr.Number(label="Score (0 - 100)"),
         gr.Number(label="Current ATMs"),
-        # gr.Number(label="Ideal ATMs (Negative Binomial)"),
+        gr.Number(label="Ideal ATMs (Negative Binomial)"),
         gr.Number(label="Ideal ATMs (XGBoost)"),
         # gr.Number(label="Log Score Probability"),