Create custom_models.py

custom_models.py

@@ -0,0 +1,87 @@
+from typing import Optional
+from transformers import PreTrainedModel, PretrainedConfig, DistilBertModel, BertModel
+import torch
+from torch import nn
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+class TransformerBasedModelDistilBert(nn.Module):
+    def __init__(self):
+        super(TransformerBasedModelDistilBert, self).__init__()
+        self.bert = DistilBertModel.from_pretrained('distilbert-base-uncased')
+        self.dropout = nn.Dropout(0.55)
+        self.fc = nn.Linear(768, 2)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+        input_shape = input_ids.size()
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+            
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.last_hidden_state[:, 0, :]
+        pooled_output = self.dropout(pooled_output)
+        logits = self.fc(pooled_output)
+        return logits
+
+class TransformerBasedModelBert(nn.Module):
+    def __init__(self):
+        super(TransformerBasedModelBert, self).__init__()
+        self.bert = BertModel.from_pretrained('bert-base-uncased')
+        self.dropout = nn.Dropout(0.55)
+        self.fc = nn.Linear(768, 2)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+        input_shape = input_ids.size()
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+            
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs[1]
+        pooled_output = self.dropout(pooled_output)
+        logits = self.fc(pooled_output)
+        return logits
+
+class MyConfigDistil(PretrainedConfig):
+    model_type = "distilbert"
+    def __init__(self, final_dropout=0.55, **kwargs):
+        super().__init__(**kwargs)
+        self.final_dropout = final_dropout
+
+class MyConfig(PretrainedConfig):
+    model_type = "bert"
+    def __init__(self, final_dropout=0.55, **kwargs):
+        super().__init__(**kwargs)
+        self.final_dropout = final_dropout
+        
+class MyHFModel_DistilBertBased(PreTrainedModel):
+    config_class = MyConfigDistil
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = TransformerBasedModelDistilBert()
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+        input_shape = input_ids.size()
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+
+        return self.model(input_ids=input_ids, attention_mask=attention_mask)
+
+class MyHFModel_BertBased(PreTrainedModel):
+    config_class = MyConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = TransformerBasedModelBert()
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+        input_shape = input_ids.size()
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+
+        return self.model(input_ids=input_ids, attention_mask=attention_mask) 
+
+config = MyConfigDistil(0.55)
+HF_DistilBertBasedModelAppDocs = MyHFModel_DistilBertBased(config)
+
+config_db = MyConfig(0.55)
+HF_BertBasedModelAppDocs = MyHFModel_BertBased(config_db)