dockerfile

BERTmodel.py

@@ -1,31 +1,15 @@
-from transformers import BertTokenizer, BertModel
 import torch
-import matplotlib.pyplot as plt
 import torch.nn as nn
-from transformers import BertTokenizer, BertModel, DataCollatorForLanguageModeling
-from datasets import load_dataset
-from torch.utils.data import DataLoader
-import torch.nn.functional as F
-
-from transformers import (
-    BertTokenizer, BertModel,
-    DataCollatorForLanguageModeling
-)
-import torch.optim as optim
+from transformers import BertTokenizer
 
-import os
-from transformers.models.bert.modeling_bert import BertOnlyMLMHead
 from models import TransformerVisualizer
 
 from transformers import (
     BertTokenizer,
     BertForMaskedLM,
     BertForSequenceClassification,
-    BertForQuestionAnswering,
 )
-import torch
 import torch.nn.functional as F
-from models import TransformerVisualizer
 
  
     

DISTILLBERTmodel.py

@@ -1,8 +1,4 @@
 import torch
-import matplotlib.pyplot as plt
-import torch.nn as nn
-from datasets import load_dataset
-from torch.utils.data import DataLoader
 import torch.nn.functional as F
 
  

ROBERTAmodel.py

@@ -3,7 +3,7 @@ import torch
 import torch.nn.functional as F
 from models import TransformerVisualizer
 from transformers import (
-    RobertaForMaskedLM, RobertaForSequenceClassification, RobertaForQuestionAnswering,
+    RobertaForMaskedLM, RobertaForSequenceClassification
 )
 
 class RoBERTaVisualizer(TransformerVisualizer):