Add HuggingFace compatible model files for bert_ffnn

config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b2b9861148aff6e66b6d3e2e1c7a8f688dd68b5a92489471713bc9aac0a6431b
-size 263
+oid sha256:f241b26cd5a3aeeeaf9f4412255776a7a578ad0f8f9174fa2fdafc61651c384f
+size 285

configuration_bert_ffnn.py

@@ -0,0 +1,27 @@
+from transformers import PretrainedConfig
+
+class BertFFNNConfig(PretrainedConfig):
+    model_type = "bert_ffnn"
+
+    def __init__(
+        self,
+        bert_model_name="microsoft/deberta-v3-base",
+        hidden_dims=[192, 96],
+        output_dim=5,
+        dropout=0.2,
+        pooling="attention",
+        freeze_bert=False,
+        freeze_layers=0,
+        use_layer_norm=True,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.bert_model_name = bert_model_name
+        self.hidden_dims = hidden_dims
+        self.output_dim = output_dim
+        self.dropout = dropout
+        self.pooling = pooling
+        self.freeze_bert = freeze_bert
+        self.freeze_layers = freeze_layers
+        self.use_layer_norm = use_layer_norm

modeling_bert_ffnn.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, AutoModel
+from .configuration_bert_ffnn import BertFFNNConfig
+
+
+class AttentionPooling(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.attention = nn.Linear(hidden_size, 1)
+
+    def forward(self, hidden_states, attention_mask):
+        scores = self.attention(hidden_states).squeeze(-1)
+        scores = scores.masked_fill(attention_mask == 0, -1e9)
+        weights = torch.softmax(scores, dim=-1)
+        return torch.sum(hidden_states * weights.unsqueeze(-1), dim=1)
+
+
+class BERT_FFNN(PreTrainedModel):
+    config_class = BertFFNNConfig
+    base_model_prefix = "bert_ffnn"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = AutoModel.from_pretrained(config.bert_model_name)
+        self.pooling = config.pooling
+        self.use_layer_norm = config.use_layer_norm
+
+        if self.pooling == "attention":
+            self.attention_pool = AttentionPooling(self.bert.config.hidden_size)
+        if config.freeze_bert:
+            for p in self.bert.parameters():
+                p.requires_grad = False
+        elif config.freeze_layers > 0:
+            for layer in self.bert.encoder.layer[:config.freeze_layers]:
+                for p in layer.parameters():
+                    p.requires_grad = False
+
+        layers = []
+        in_dim = self.bert.config.hidden_size
+        for h_dim in config.hidden_dims:
+            layers.append(nn.Linear(in_dim, h_dim))
+            layers.append(nn.ReLU())
+            if config.use_layer_norm:
+                layers.append(nn.LayerNorm(h_dim))
+            layers.append(nn.Dropout(config.dropout))
+            in_dim = h_dim
+
+        layers.append(nn.Linear(in_dim, config.output_dim))
+        self.classifier = nn.Sequential(*layers)
+
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+
+        if self.pooling == "mean":
+            mask = attention_mask.unsqueeze(-1).float()
+            sum_emb = (outputs.last_hidden_state * mask).sum(1)
+            features = sum_emb / mask.sum(1).clamp(min=1e-9)
+        elif self.pooling == "max":
+            mask = attention_mask.unsqueeze(-1).float()
+            masked_emb = outputs.last_hidden_state.masked_fill(mask == 0, float('-inf'))
+            features, _ = masked_emb.max(dim=1)
+        elif self.pooling == "attention":
+            features = self.attention_pool(outputs.last_hidden_state, attention_mask)
+        else:  # CLS pooling
+            features = (
+                outputs.pooler_output
+                if getattr(outputs, "pooler_output", None) is not None
+                else outputs.last_hidden_state[:, 0]
+            )
+
+        logits = self.classifier(features)
+        return logits