Update app.py

app.py

@@ -1,19 +1,35 @@
 import gradio as gr
 import torch
 import numpy as np
+import pandas as pd
 
-# ======== Example Dummy Model ========
-# Replace this with your own model loading
-class DummyModel(torch.nn.Module):
-    def __init__(self):
+dataset = pd.read_csv("cleaned_df_noindex.csv")
+df_amenities = pd.read_csv("df_amenities.csv")
+
+class NegBinomialModel(nn.Module):
+    def __init__(self, in_features):
         super().__init__()
-        self.linear = torch.nn.Linear(2, 3)  # 2 inputs (lat, lon) → 3 outputs
+        self.linear = nn.Linear(in_features, 1)
+        self.log_alpha = nn.Parameter(torch.tensor(0.0))
 
     def forward(self, x):
-        return self.linear(x)
+        eta = self.linear(x)
+        mu = torch.exp(eta).squeeze(-1)
+        alpha = torch.exp(self.log_alpha)
+        return mu, alpha
 
-# load your trained model (replace with torch.load if you have .pt file)
-model = DummyModel()
+def negbinom_loss(y, mu, alpha):
+    log_prob = (
+        torch.lgamma(y + 1/alpha)
+        - torch.lgamma(1/alpha)
+        - torch.lgamma(y + 1)
+        + (1/alpha) * torch.log(1 / (1 + alpha * mu))
+        + y * torch.log((alpha * mu) / (1 + alpha * mu))
+    )
+    return -torch.mean(log_prob)
+    
+model = NegBinomialModel(len(dataset.columns)-2)
+model.load_state_dict(torch.load("model_weights.pt"))
 model.eval()
 
 # ======== Prediction Function ========
@@ -26,13 +42,15 @@ def predict_score(lat, lon):
         outputs = model(inputs).numpy().flatten()
 
     # Unpack into respective values
-    score, num_banks, normal_score = outputs
+    mu_pred, alpha = outputs
+
+    score = (1 * np.abs(mu_pred + 0.1)) * 100
 
     # You can apply any post-processing here
     return {
         "Score": round(float(score), 3),
-        "Num Banks": round(float(num_banks), 3),
-        "Normal Score": round(float(normal_score), 3),
+        "Num Banks": round(float(mu_pred), 3),
+        # "Normal Score": round(float(normal_score), 3),
     }
 
 # ======== Gradio Interface ========