Upload probe.py with huggingface_hub

probe.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+from typing import Literal
+from transformers import PretrainedConfig, PreTrainedModel
+
+
+class ProbeConfig(PretrainedConfig):
+    model_type = "linear_probe"
+
+    def __init__(
+        self,
+        embedding_dim: int = 768,
+        dropout: float = 0.0,
+        layer_index: int = -1,
+        probe_type: Literal["linear", "nonlinear"] = "linear",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.embedding_dim = embedding_dim
+        self.dropout = dropout
+        self.layer_index = layer_index
+        self.probe_type = probe_type
+
+
+class ProbeModel(PreTrainedModel):
+    config_class = ProbeConfig
+
+    def __init__(self, config: ProbeConfig):
+        super().__init__(config)
+        self.dropout = nn.Dropout(config.dropout) if config.dropout > 0 else None
+        self.linear = nn.Linear(config.embedding_dim, 1)
+
+    def forward(
+        self,
+        embeddings: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        if self.dropout is not None:
+            embeddings = self.dropout(embeddings)
+        logits = self.linear(embeddings)
+        return torch.sigmoid(logits)