Remove warning about attn_implementation for new transformers versions (#63)

Author: avaucher

rxnmapper/core.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+from contextlib import contextmanager
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import numpy as np
@@ -28,6 +29,17 @@
 _logger.addHandler(logging.NullHandler())
 
 
+@contextmanager
+def suppress_transformers_warnings():
+    logger = logging.getLogger("transformers")
+    previous_level = logger.level
+    logger.setLevel(logging.ERROR)
+    try:
+        yield
+    finally:
+        logger.setLevel(previous_level)
+
+
 class RXNMapper:
     """Wrap the Transformer model, corresponding tokenizer, and attention scoring algorithms.
 
@@ -136,8 +148,11 @@ def convert_batch_to_attns(
                 f"Reaction SMILES has {max_input_length} tokens, should be at most {max_supported_by_model}."
             )
 
-        with torch.no_grad():
-            output = self.model(**parsed_input)
+        # suppress warning that suggests setting "attn_implementation"; doing
+        # so would break compatibility for old `transformers` versions.
+        with suppress_transformers_warnings():
+            with torch.no_grad():
+                output = self.model(**parsed_input)
         attentions = output[2]
         selected_attns = torch.cat(
             [a.unsqueeze(1) for i, a in enumerate(attentions) if i in use_layers],