Update model_archs

model_archs

@@ -117,3 +117,28 @@ class ResNet1d(torch.nn.Module):
     def forward(self, x):
         return self.blocks(x)
 
+
+
+class LogisticRegressionTorch(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int):
+        super(LogisticRegressionTorch, self).__init__()
+        self.batch_norm = nn.BatchNorm1d(num_features=input_dim)
+        self.linear = nn.Linear(input_dim, output_dim)
+
+    def forward(self, x):
+        x = self.batch_norm(x)
+        out = self.linear(x)
+        return out
+
+class BertClassifier(nn.Module):
+    def __init__(self, bert_model: AutoModel, classifier: LogisticRegressionTorch, num_labels: int):
+        super(BertClassifier, self).__init__()
+        self.bert = bert_model
+        self.classifier = classifier
+        self.num_labels = num_labels
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor = None):
+        outputs = self.bert(input_ids, attention_mask=attention_mask, output_hidden_states=True)
+        pooled_output = outputs.hidden_states[-1][:, 0, :]
+        logits = self.classifier(pooled_output)
+        return logits