Update models.py

models.py

@@ -2,9 +2,21 @@ import torch
 from transformers import AutoModelForTokenClassification, AutoModelForSequenceClassification, AdamW, get_linear_schedule_with_warmup
 from transformers import BertForTokenClassification, BertForSequenceClassification,BertPreTrainedModel, BertModel
 import torch.nn as nn
-from .utils import *
 import torch.nn.functional as F
 
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
 
 
 class Model_Rational_Label(BertPreTrainedModel):