| import torch |
| import torch.nn.functional as F |
| import json |
| import os |
| from huggingface_hub import hf_hub_download |
| from model import PhysicsGuidedBNN |
|
|
| class PGBNNPredictor: |
| """ |
| Inference and Uncertainty Quantification engine. |
| """ |
| def __init__(self, model: PhysicsGuidedBNN, config: dict): |
| self.model = model |
| self.config = config |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.model.to(self.device) |
| self.model.eval() |
|
|
| @classmethod |
| def from_pretrained(cls, repo_id: str): |
| config_path = hf_hub_download(repo_id=repo_id, filename="config.json") |
| weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin") |
|
|
| with open(config_path, "r") as f: |
| config = json.load(f) |
|
|
| model = PhysicsGuidedBNN(config) |
| model.load_state_dict(torch.load(weights_path, map_location="cpu")) |
| return cls(model, config) |
|
|
| def predict(self, x: torch.Tensor, num_samples: int = 30): |
| """ |
| Performs Monte Carlo forward passes to compute predictive mean and epistemic uncertainty. |
| """ |
| x = x.to(self.device) |
| probs_list = [] |
|
|
| with torch.no_grad(): |
| for _ in range(num_samples): |
| logits = self.model(x) |
| probs = F.softmax(logits, dim=-1)[:, 1] |
| probs_list.append(probs.unsqueeze(0)) |
|
|
| probs_tensor = torch.cat(probs_list, dim=0) |
| |
| |
| mean_prob = torch.mean(probs_tensor, dim=0) |
| |
| |
| epistemic_uncertainty = torch.std(probs_tensor, dim=0) |
|
|
| return mean_prob, epistemic_uncertainty |
