Create modeling_custom.py

modeling_custom.py

@@ -0,0 +1,29 @@
+from torch import nn
+from huggingface_hub import PyTorchModelHubMixin
+
+class GatedRegressionModel(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, base_model):
+        super().__init__()
+        self.base_model = base_model
+        self.hidden_dim = 1024
+        self.gate_model = nn.ModuleList([
+            nn.Linear(self.base_model.pre_classifier.out_features, self.hidden_dim, bias=True),
+            nn.GELU(),
+            nn.Dropout(p=0.2),
+            nn.Linear(self.hidden_dim, self.hidden_dim, bias=True),
+            nn.GELU(),
+            nn.Dropout(p=0.2),
+            nn.Linear(self.hidden_dim, 1, bias=True),
+            nn.GELU(),
+            nn.Dropout(p=0.2),
+        ])
+        
+
+    def forward(self, input_ids, attention_mask):
+        o = self.base_model.distilbert(input_ids=input_ids, attention_mask=attention_mask)[0][:, 0] # encoder-decoder architecture
+        o = self.base_model.pre_classifier(o)
+        scores = self.base_model.classifier(o)
+        gate = o
+        for layer in self.gate_model:
+            gate = layer(gate)
+        return (gate * scores).sum(axis=1)