fix bt test failures due to default sdpa attention

huggingface · May 23, 2024 · c9deaec · c9deaec
1 parent e0f5812
commit c9deaec
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/optimum/pipelines/pipelines_base.py b/optimum/pipelines/pipelines_base.py
@@ -179,7 +179,7 @@ def load_bettertransformer(
     **kwargs,
 ):
     if model_kwargs is None:
-        model_kwargs = {}
+        model_kwargs = {"attn_implementation": "eager"}
 
     if model is None:
         model_id = SUPPORTED_TASKS[targeted_task]["default"]

diff --git a/tests/bettertransformer/test_encoder.py b/tests/bettertransformer/test_encoder.py
@@ -114,7 +114,7 @@ def test_inference_speed(self):
         """
         model_name = "bert-base-uncased"
 
-        hf_model = AutoModel.from_pretrained(model_name).eval()
+        hf_model = AutoModel.from_pretrained(model_name, attn_implementation="eager").eval()
         bt_model = BetterTransformer.transform(hf_model, keep_original_model=True)
 
         BATCH_SIZE = 8

diff --git a/tests/bettertransformer/testing_utils.py b/tests/bettertransformer/testing_utils.py
@@ -235,7 +235,7 @@ def _test_logits(self, model_id: str, model_type: str, **preprocessor_kwargs):
         inputs = self.prepare_inputs_for_class(model_id=model_id, model_type=model_type, **preprocessor_kwargs)
 
         torch.manual_seed(0)
-        hf_random_model = AutoModel.from_pretrained(model_id).eval()
+        hf_random_model = AutoModel.from_pretrained(model_id, attn_implementation="eager").eval()
         random_config = hf_random_model.config
 
         hf_random_model = hf_random_model.eval()