feat: adding secure loading of models by default for haystack

haystack/__init__.py

@@ -20,6 +20,8 @@
 from haystack.schema import Document, Answer, Label, MultiLabel, Span, EvaluationResult
 from haystack.nodes.base import BaseComponent
 from haystack.pipelines.base import Pipeline
+from haystack.environment import set_pytorch_secure_model_loading
 
 
 pd.options.display.max_colwidth = 80
+set_pytorch_secure_model_loading()

haystack/environment.py

@@ -18,6 +18,12 @@
 env_meta_data: Dict[str, Any] = {}
 
 
+def set_pytorch_secure_model_loading(flag_val="1"):
+    # To load secure only model pytorch requires value of
+    # TORCH_FORCE_WEIGHTS_ONLY_LOAD to be ["1", "y", "yes", "true"]
+    os.environ["TORCH_FORCE_WEIGHTS_ONLY_LOAD"] = flag_val
+
+
 def get_or_create_env_meta_data() -> Dict[str, Any]:
     """
     Collects meta data about the setup that is used with Haystack, such as: operating system, python version, Haystack version, transformers version, pytorch version, number of GPUs, execution environment, and the value stored in the env variable HAYSTACK_EXECUTION_CONTEXT.

test/others/test_utils.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from random import random
 from typing import List
 
@@ -1245,6 +1246,21 @@ def greet2(name: str):
     assert greet2("John") == "Hello John"
 
 
+def test_secure_model_loading():
+    # setting the flag explicitly to zero
+    os.environ["TORCH_FORCE_WEIGHTS_ONLY_LOAD"] = "0"
+    env_val = os.getenv("TORCH_FORCE_WEIGHTS_ONLY_LOAD")
+    assert env_val == "0"
+
+    # now testing if just importing haystack is enough to enable secure loading of pytorch models
+    import importlib
+    import haystack
+
+    importlib.reload(haystack)
+    env_val = os.getenv("TORCH_FORCE_WEIGHTS_ONLY_LOAD", "0")
+    assert env_val == "1"
+
+
 class TestAggregateLabels:
     @pytest.fixture
     def standard_labels(self) -> List[Label]: