probe_clf.py

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, device="cpu"):
        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),
        ).to(device)
        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)