Update app.py

app.py

@@ -12,23 +12,14 @@ class NegBinomialModel(nn.Module):
     def __init__(self, in_features):
         super().__init__()
         self.linear = nn.Linear(in_features, 1)
-        self.log_alpha = nn.Parameter(torch.tensor(0.0))
+        self.alpha = nn.Parameter(torch.tensor(0.5))
 
     def forward(self, x):
-        eta = self.linear(x)
-        mu = torch.exp(eta).squeeze(-1)
-        alpha = torch.exp(self.log_alpha)
-        return mu, alpha
+        # 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
 
-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(17)
 model.load_state_dict(torch.load("model_weights.pt", map_location='cpu'))