fix pickle unsafe-load

paddlenlp/utils/serialization.py

@@ -160,6 +160,24 @@
         return super().find_class(mod_name, name)
 
 
+class SafeUnpickler(pickle.Unpickler):
+    """
+    A safe unpickler that only allows loading of built-in basic data types.
+    """
+
+    def find_class(self, module, name):
+        """
+        Overrides the find_class method to only allow loading of built-in basic data types.
+
+        :param module: The module name.
+        :param name: The class name.
+        :return: The class if allowed, otherwise raises UnpicklingError.
+        """
+        if module == "builtins" and name in {"int", "float", "str", "tuple", "list", "dict", "set"}:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(f"Unsafe object loading is prohibited: {module}.{name}")
+
+
 def _rebuild_tensor_stage(storage, storage_offset, size, stride, requires_grad, backward_hooks):
     # if a tensor has shape [M, N] and stride is [1, N], it's column-wise / fortran-style
     # if a tensor has shape [M, N] and stride is [M, 1], it's row-wise / C-style

slm/model_zoo/chinesebert/utils.py

@@ -29,6 +29,7 @@
     LinearDecayWithWarmup,
     PolyDecayWithWarmup,
 )
+from paddlenlp.utils.serialization import SafeUnpickler
 
 scheduler_type2cls = {
     "linear": LinearDecayWithWarmup,
@@ -121,7 +122,7 @@ def save_pickle(data, file_path):
 
 def load_pickle(input_file):
     with open(str(input_file), "rb") as f:
-        data = pickle.load(f)
+        data = SafeUnpickler(f).load()
     return data
 
 

slm/model_zoo/t5/utils.py

@@ -28,6 +28,7 @@
     LinearDecayWithWarmup,
     PolyDecayWithWarmup,
 )
+from paddlenlp.utils.serialization import SafeUnpickler
 
 
 def accuracy(targets, predictions):
@@ -158,5 +159,5 @@ def save_pickle(data, file_path):
 
 def load_pickle(input_file):
     with open(str(input_file), "rb") as f:
-        data = pickle.load(f)
+        data = SafeUnpickler(f).load()
     return data