Create pisco_clf.py

pisco_clf.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class PISCOClassifier(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, d: int, hidden: int = 512, threshold: float = 0.5):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d, hidden),
+            nn.LayerNorm(hidden),
+            nn.GELU(),
+            nn.Dropout(0.3),
+            nn.Linear(hidden, hidden // 4),
+            nn.GELU(),
+            nn.Dropout(0.2),
+            nn.Linear(hidden // 4, 1),
+        )
+        self.threshold = threshold
+
+    def forward(self, x):
+        return self.net(x).squeeze(-1)
+
+    @torch.inference_mode()
+    def predict_proba(self, X) -> np.ndarray:
+        self.eval()
+        x = self._as_tensor(X)
+        return torch.sigmoid(self.net(x)).cpu().numpy()
+
+    def predict(self, X, threshold: float | None = None) -> np.ndarray:
+        """Binary predictions. Uses stored threshold if not given."""
+        t = threshold if threshold is not None else self.threshold
+        return (self.predict_proba(X) >= t).astype(int)
+
+    @staticmethod
+    def _as_tensor(X) -> torch.Tensor:
+        if isinstance(X, torch.Tensor):
+            return X.float()
+        return torch.tensor(np.asarray(X), dtype=torch.float32)