remove extra imports

.flake8

@@ -0,0 +1,2 @@
+[flake8]
+max-line-length = 127

examples/bert_snli.py

@@ -1,8 +1,6 @@
 import argparse
-import collections
 import logging
 import os
-import sys
 
 import torch
 import transformers
@@ -35,7 +33,8 @@ def main():
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
     tokenizer = transformers.AutoTokenizer.from_pretrained(args.base_model, cache_dir=args.model_cache_dir)
-    model = transformers.AutoModelForSequenceClassification.from_pretrained(args.base_model, num_labels=3, cache_dir=args.model_cache_dir)
+    model = transformers.AutoModelForSequenceClassification.from_pretrained(
+        args.base_model, num_labels=3, cache_dir=args.model_cache_dir)
     model.to(device)
     swag_model = SwagBertForSequenceClassification.from_base(model)
     swag_model.to(device)

examples/load_pretrained.py

@@ -1,13 +1,9 @@
 import argparse
-import collections
 import logging
-import os
-import sys
 
 import transformers
 
 from swag_transformers.swag_bert import SwagBertConfig, SwagBertForSequenceClassification
-from swag_transformers.trainer_utils import SwagUpdateCallback
 
 
 def main():

examples/marian_mt.py

@@ -1,8 +1,6 @@
 import argparse
-import collections
 import logging
 import os
-import sys
 
 import torch
 import transformers

src/swag_transformers/base.py

@@ -1,10 +1,9 @@
 """SWAG wrapper base classes"""
 
-from abc import ABCMeta, abstractmethod
 import copy
 import functools
 import logging
-from typing import Union, Type
+from typing import Type
 
 import torch
 

src/swag_transformers/trainer_utils.py

@@ -2,7 +2,6 @@
 
 import logging
 
-from swag.posteriors.swag import SWAG
 from transformers import TrainerCallback
 
 

tests/test_swag_bart.py

@@ -2,16 +2,12 @@
 import unittest
 import tempfile
 
-import numpy as np
 import torch
 
-from datasets import Dataset, DatasetDict
-from transformers import AutoConfig, AutoModel, AutoModelForSequenceClassification, \
-    AutoTokenizer, DataCollatorWithPadding, Trainer, TrainingArguments, BartForConditionalGeneration
+from transformers import AutoTokenizer, BartForConditionalGeneration
 
 from swag_transformers.swag_bart import SwagBartConfig, SwagBartModel, SwagBartPreTrainedModel, \
-    SwagBartForSequenceClassification, SwagBartForConditionalGeneration
-from swag_transformers.trainer_utils import SwagUpdateCallback
+    SwagBartForConditionalGeneration
 
 
 class TestSwagBart(unittest.TestCase):
@@ -77,7 +73,6 @@ def test_pretrained_bart_generative(self):
         self.assertEqual(base_out, out)
 
 
-
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
     unittest.main()

tests/test_swag_bert.py

@@ -2,11 +2,10 @@
 import unittest
 import tempfile
 
-import numpy as np
 import torch
 
 from datasets import Dataset, DatasetDict
-from transformers import AutoConfig, AutoModel, AutoModelForSequenceClassification, AutoModelWithLMHead, \
+from transformers import AutoModel, AutoModelForSequenceClassification, \
     AutoTokenizer, DataCollatorWithPadding, Trainer, TrainingArguments
 
 from swag_transformers.swag_bert import SwagBertConfig, SwagBertLMHeadModel, SwagBertModel, SwagBertPreTrainedModel, \