model replaced with large

model.py

@@ -17,8 +17,8 @@ config = dict(
     num_labels=2,
     
     # model info
-    tokenizer_path = 'allenai/biomed_roberta_base', # 'roberta-base', 
-    model_checkpoint = 'allenai/biomed_roberta_base', # 'roberta-base', 
+    tokenizer_path = 'roberta-large', # 'allenai/biomed_roberta_base',
+    model_checkpoint = 'model_large_pseudo_label.pth', # 'allenai/biomed_roberta_base', 
     device = 'cuda' if torch.cuda.is_available() else 'cpu',
     
     # training paramters
@@ -78,22 +78,106 @@ class NBMETestData(torch.utils.data.Dataset):
             'sequence_ids': sequence_ids,
         }
 
+# class NBMEModel(nn.Module):
+#     def __init__(self, num_labels=1, path=None):
+#         super().__init__()
+        
+#         layer_norm_eps: float = 1e-6
+        
+#         self.path = path
+#         self.num_labels = num_labels
+        
+#         self.transformer = transformers.AutoModel.from_pretrained(config['model_checkpoint'])
+#         self.dropout = nn.Dropout(0.2)
+#         self.output = nn.Linear(768, 1)
+        
+#         if self.path is not None:
+#             self.load_state_dict(torch.load(self.path)['model'])
+
+#     def forward(self, data):
+        
+#         ids = data['input_ids']
+#         mask = data['attention_mask']
+#         try:
+#             target = data['targets']
+#         except:
+#             target = None
+
+#         transformer_out = self.transformer(ids, mask)
+#         sequence_output = transformer_out[0]
+#         sequence_output = self.dropout(sequence_output)
+#         logits = self.output(sequence_output)
+
+#         ret = {
+#             "logits": torch.sigmoid(logits),
+#         }
+        
+#         if target is not None:
+#             loss = self.get_loss(logits, target)
+#             ret['loss'] = loss
+#             ret['targets'] = target
+
+#         return ret
+
+        
+#     def get_optimizer(self, learning_rate, weigth_decay):
+#         optimizer = torch.optim.AdamW(
+#             self.parameters(), 
+#             lr=learning_rate, 
+#             weight_decay=weigth_decay,
+#         )
+#         if self.path is not None:
+#             optimizer.load_state_dict(torch.load(self.path)['optimizer'])
+        
+#         return optimizer
+            
+#     def get_scheduler(self, optimizer, num_warmup_steps, num_training_steps):
+#         scheduler = transformers.get_linear_schedule_with_warmup(
+#             optimizer,
+#             num_warmup_steps=num_warmup_steps,
+#             num_training_steps=num_training_steps,
+#         )
+#         if self.path is not None:
+#             scheduler.load_state_dict(torch.load(self.path)['scheduler'])
+            
+#         return scheduler
+    
+#     def get_loss(self, output, target):
+#         loss_fn = nn.BCEWithLogitsLoss(reduction="none")
+#         loss = loss_fn(output.view(-1, 1), target.view(-1, 1))
+#         loss = torch.masked_select(loss, target.view(-1, 1) != -100).mean()
+#         return loss
+
+
 class NBMEModel(nn.Module):
-    def __init__(self, num_labels=1, path=None):
+    def __init__(self, num_labels=2, path=None):
         super().__init__()
         
         layer_norm_eps: float = 1e-6
         
         self.path = path
         self.num_labels = num_labels
+        self.config = transformers.AutoConfig.from_pretrained(config['model_checkpoint'])
+
+        self.config.update(
+            {
+                "layer_norm_eps": layer_norm_eps,
+            }
+        )
+        self.transformer = transformers.AutoModel.from_pretrained(config['model_checkpoint'], config=self.config)
+        self.dropout = nn.Dropout(0.1)
         
-        self.transformer = transformers.AutoModel.from_pretrained(config['model_checkpoint'])
-        self.dropout = nn.Dropout(0.2)
-        self.output = nn.Linear(768, 1)
+        self.dropout1 = nn.Dropout(0.1)
+        self.dropout2 = nn.Dropout(0.2)
+        self.dropout3 = nn.Dropout(0.3)
+        self.dropout4 = nn.Dropout(0.4)
+        self.dropout5 = nn.Dropout(0.5)
+        
+        self.output = nn.Linear(self.config.hidden_size, 1)
         
         if self.path is not None:
             self.load_state_dict(torch.load(self.path)['model'])
-
+    
     def forward(self, data):
         
         ids = data['input_ids']
@@ -106,16 +190,29 @@ class NBMEModel(nn.Module):
         transformer_out = self.transformer(ids, mask)
         sequence_output = transformer_out[0]
         sequence_output = self.dropout(sequence_output)
-        logits = self.output(sequence_output)
+    
+        logits1 = self.output(self.dropout1(sequence_output))
+        logits2 = self.output(self.dropout2(sequence_output))
+        logits3 = self.output(self.dropout3(sequence_output))
+        logits4 = self.output(self.dropout4(sequence_output))
+        logits5 = self.output(self.dropout5(sequence_output))
 
+        logits = (logits1 + logits2 + logits3 + logits4 + logits5) / 5
         ret = {
-            "logits": torch.sigmoid(logits),
+            'logits': torch.sigmoid(logits), 
         }
         
+        loss = 0
+
         if target is not None:
-            loss = self.get_loss(logits, target)
+            loss1 = self.get_loss(logits1, target)
+            loss2 = self.get_loss(logits2, target)
+            loss3 = self.get_loss(logits3, target)
+            loss4 = self.get_loss(logits4, target)
+            loss5 = self.get_loss(logits5, target)
+            loss = (loss1 + loss2 + loss3 + loss4 + loss5) / 5
             ret['loss'] = loss
-            ret['targets'] = target
+            ret['target'] = target
 
         return ret
 
@@ -148,6 +245,7 @@ class NBMEModel(nn.Module):
         loss = torch.masked_select(loss, target.view(-1, 1) != -100).mean()
         return loss
 
+
 def get_location_predictions(preds, offset_mapping, sequence_ids, test=False):
     all_predictions = []
     for pred, offsets, seq_ids in zip(preds, offset_mapping, sequence_ids):