improve test code and transformers support (#5)

* add setting random seed to examples
* add support for transformers >= 4.43
* fix tests

Author: svirpioj

examples/bert_snli.py

@@ -27,13 +27,17 @@ def main():
                         help="number of steps between collecting parameters; set to zero for per epoch updates")
     parser.add_argument("--learning-rate", type=float, default=2e-5, help="learning rate")
     parser.add_argument("--swag-modules", type=str, action='append', help="restrict SWAG to modules matching given prefix(es)")
+    parser.add_argument("--seed", type=int, default=None, help="set random seed")
     args = parser.parse_args()
 
     if args.device:
         device = args.device
     else:
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    if args.seed is not None:
+        transformers.set_seed(args.seed)
+
     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)

examples/marian_mt.py

@@ -24,13 +24,17 @@ def main():
     parser.add_argument("--collect-steps", type=int, default=100, help="number of steps between collecting parameters")
     parser.add_argument("--learning-rate", type=float, default=2e-5, help="learning rate")
     parser.add_argument("--swag-modules", type=str, action='append', help="restrict SWAG to modules matching given prefix(es)")
+    parser.add_argument("--seed", type=int, default=None, help="set random seed")
     args = parser.parse_args()
 
     if args.device:
         device = args.device
     else:
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    if args.seed is not None:
+        transformers.set_seed(args.seed)
+
     tokenizer = transformers.AutoTokenizer.from_pretrained(args.base_model)
     model = transformers.MarianMTModel.from_pretrained(args.base_model)
     model.to(device)

setup.py

@@ -18,9 +18,8 @@
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     install_requires=[
-        "transformers>=4.30",
-        "transformers[torch]>=4.30,<4.43",
-        "swa_gaussian>=0.1.8"
+        "transformers[torch]>=4.30",
+        "swa_gaussian>=0.1.9"
     ],
     extras_require={
         "test": ["datasets", "pytest", "sentencepiece"]

tests/test_swag_bart.py

@@ -13,8 +13,8 @@
 
 class TestSwagBart(unittest.TestCase):
 
-    pretrained_model_name = 'Finnish-NLP/bart-small-finnish'
-    # pretrained_model_name = 'sshleifer/bart-tiny-random'
+    # pretrained_model_name = 'Finnish-NLP/bart-small-finnish'
+    pretrained_model_name = 'sshleifer/bart-tiny-random'
 
     def test_untrained(self):
         hidden_size = 240
@@ -45,12 +45,11 @@ def pretrained_bart_generative(self, no_cov_mat):
         tokenizer = AutoTokenizer.from_pretrained(self.pretrained_model_name, clean_up_tokenization_spaces=False)
 
         gen_config = GenerationConfig.from_model_config(model.config)
-        logging.warning(gen_config)
         gen_config.max_new_tokens = 10
-        logging.warning(gen_config)
+        logging.debug(gen_config)
 
         swag_model.swag.collect_model(model)
-        swag_model.sample_parameters(cov=not no_cov_mat)
+        swag_model.sample_parameters(cov=not no_cov_mat, seed=1234)
         # has to be updated manually when using collect_model directly
         swag_model.config.cov_mat_rank = swag_model.swag.cov_mat_rank
 
@@ -61,29 +60,26 @@ def pretrained_bart_generative(self, no_cov_mat):
 
         # Test generate
         example = "I have no BART and I must generate"
-        torch.manual_seed(123)
         batch = tokenizer(example, return_tensors="pt")
         base_generated_ids = model.generate(batch["input_ids"], generation_config=gen_config)
-        # max_length=20, num_beams=1, do_sample=False, early_stopping=False
         base_out = tokenizer.batch_decode(base_generated_ids, skip_special_tokens=True)
-        logging.warning(base_out)
 
         generated_ids = swag_model.generate(batch["input_ids"], generation_config=gen_config)
         out = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-        logging.warning(out)
         self.assertEqual(base_out, out)
 
         # Test saving & loading
         with tempfile.TemporaryDirectory() as tempdir:
             swag_model.save_pretrained(tempdir)
-            logging.warning(os.listdir(tempdir))
-            with open(os.path.join(tempdir, 'config.json'), 'r') as fobj:
-                logging.warning(fobj.read())
+            logging.debug(os.listdir(tempdir))
+            with open(os.path.join(tempdir, 'config.json'), 'r', encoding='utf8') as fobj:
+                logging.debug(fobj.read())
             stored_model = SwagBartForConditionalGeneration.from_pretrained(tempdir).to(device)
 
+        stored_model.sample_parameters(cov=not no_cov_mat, seed=1234)
         stored_fwd_out = stored_model.forward(
             input_ids=torch.tensor([[3, 14]]), decoder_input_ids=torch.tensor([[1, 2, 4]]))
-        self.assertTrue(torch.allclose(swag_fwd_out.logits, stored_fwd_out.logits))
+        self.assertTrue(torch.allclose(swag_fwd_out.logits, stored_fwd_out.logits, atol=1e-06))
 
         generated_ids = stored_model.generate(batch["input_ids"], generation_config=gen_config)
         out = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

tests/test_swag_bert.py

@@ -203,7 +203,7 @@ def tokenize_function(example):
                 training_args,
                 train_dataset=tokenized_datasets["train"],
                 data_collator=data_collator,
-                tokenizer=tokenizer,
+                processing_class=tokenizer,
                 callbacks=[SwagUpdateCallback(swag_model)]
             )
             trainer.train()

tests/test_swag_marian.py

@@ -185,7 +185,7 @@ def tokenize_function(example):
                 training_args,
                 train_dataset=tokenized_datasets["train"],
                 data_collator=data_collator,
-                tokenizer=tokenizer,
+                processing_class=tokenizer,
                 callbacks=[SwagUpdateCallback(swag_model, collect_steps=2)]
             )
             trainer.train()