Update model.py

model.py

@@ -1,39 +1,47 @@
-import torch.nn as nn
-import torch
-from transformers import PreTrainedModel, AutoModel
-from transformers import PretrainedConfig
-
-class DualDistilBERTClassifier(PreTrainedModel):
-    config_class = PretrainedConfig
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
-        self.encoder = AutoModel.from_pretrained(config.base_model_name_or_path)
-        hidden_size = self.encoder.config.hidden_size
-
-        self.classifier = nn.Sequential(
-            nn.Linear(hidden_size * 2, 512),
-            nn.ReLU(),
-            nn.Dropout(0.1),
-            nn.Linear(512, self.num_labels)
-        )
-
-    def forward(self, title_input_ids, title_attention_mask,
-                abstract_input_ids, abstract_attention_mask, labels=None):
-
-        title_out = self.encoder(input_ids=title_input_ids,
-                                 attention_mask=title_attention_mask).last_hidden_state[:, 0]
-
-        abstract_out = self.encoder(input_ids=abstract_input_ids,
-                                    attention_mask=abstract_attention_mask).last_hidden_state[:, 0]
-
-        # Если абстракт пустой — можно заменить его на нули
-        abstract_out = torch.where(
-            abstract_attention_mask.sum(dim=1).unsqueeze(1) > 0,
-            abstract_out,
-            torch.zeros_like(abstract_out)
-        )
-
-        combined = torch.cat([title_out, abstract_out], dim=1)
-        logits = self.classifier(combined)
+from transformers.utils import (
+    logging,
+    is_torch_available,
+)
+from transformers.modeling_utils import (
+    no_init_weights,
+    init_empty_weights,
+)
+import torch.nn as nn
+import torch
+from transformers import PreTrainedModel, AutoModel
+from transformers import PretrainedConfig
+
+class DualDistilBERTClassifier(PreTrainedModel):
+    config_class = PretrainedConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.encoder = AutoModel.from_pretrained(config.base_model_name_or_path)
+        hidden_size = self.encoder.config.hidden_size
+
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_size * 2, 512),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(512, self.num_labels)
+        )
+
+    def forward(self, title_input_ids, title_attention_mask,
+                abstract_input_ids, abstract_attention_mask, labels=None):
+
+        title_out = self.encoder(input_ids=title_input_ids,
+                                 attention_mask=title_attention_mask).last_hidden_state[:, 0]
+
+        abstract_out = self.encoder(input_ids=abstract_input_ids,
+                                    attention_mask=abstract_attention_mask).last_hidden_state[:, 0]
+
+        # Если абстракт пустой — можно заменить его на нули
+        abstract_out = torch.where(
+            abstract_attention_mask.sum(dim=1).unsqueeze(1) > 0,
+            abstract_out,
+            torch.zeros_like(abstract_out)
+        )
+
+        combined = torch.cat([title_out, abstract_out], dim=1)
+        logits = self.classifier(combined)
         return {'logits': logits}