From f4ba960e4137f1d8bb447eca741bb8dc50c5675e Mon Sep 17 00:00:00 2001 From: Joseph Catrambone Date: Wed, 26 Jun 2024 13:49:04 -0700 Subject: [PATCH 1/2] Bugfix: Sidestep MacOS Segfault in OpenMP. --- tests/unit_tests/mocks/mock_hf_models.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/mocks/mock_hf_models.py b/tests/unit_tests/mocks/mock_hf_models.py index 737f6f29c..e39515884 100644 --- a/tests/unit_tests/mocks/mock_hf_models.py +++ b/tests/unit_tests/mocks/mock_hf_models.py @@ -3,6 +3,9 @@ def make_mock_model_and_tokenizer(): """Returns a tuple of HF AutoModelForCausalLM and AutoTokenizer.""" + import torch + torch.set_num_threads(1) + from transformers import AutoModelForCausalLM, AutoTokenizer # Can regenerate the sample pipe with this: @@ -15,9 +18,15 @@ def make_mock_model_and_tokenizer(): os.path.abspath(os.path.normpath(os.path.dirname(__file__))), "tiny-random-gpt2" ) - model = AutoModelForCausalLM.from_pretrained(savedir, local_files_only=True) + model = AutoModelForCausalLM.from_pretrained( + savedir, + local_files_only=True, + ) - tokenizer = AutoTokenizer.from_pretrained(savedir, local_files_only=True) + tokenizer = AutoTokenizer.from_pretrained( + savedir, + local_files_only=True, + ) return model, tokenizer From 31ef63a2ed1f939ab005a6ae50d38f18afea97e7 Mon Sep 17 00:00:00 2001 From: Joseph Catrambone Date: Wed, 26 Jun 2024 13:54:07 -0700 Subject: [PATCH 2/2] Reformat. --- tests/unit_tests/mocks/mock_hf_models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/mocks/mock_hf_models.py b/tests/unit_tests/mocks/mock_hf_models.py index e39515884..e13c70349 100644 --- a/tests/unit_tests/mocks/mock_hf_models.py +++ b/tests/unit_tests/mocks/mock_hf_models.py @@ -4,6 +4,7 @@ def make_mock_model_and_tokenizer(): """Returns a tuple of HF AutoModelForCausalLM and AutoTokenizer.""" import torch + torch.set_num_threads(1) from transformers import AutoModelForCausalLM, AutoTokenizer