decompressing joblib file before checking it (NVIDIA-NeMo#13732)

Ssofja · web-flow · commit 33ea12a288cc · 2025-05-30T13:54:48.000Z
* decompressing joblib file before checking it

Signed-off-by: Ssofja &lt;sofiakostandian@gmail.com&gt;

* Apply isort and black reformatting

Signed-off-by: Ssofja &lt;Ssofja@users.noreply.github.com&gt;
Signed-off-by: Ssofja &lt;sofiakostandian@gmail.com&gt;

---------

Signed-off-by: Ssofja &lt;sofiakostandian@gmail.com&gt;
Signed-off-by: Ssofja &lt;Ssofja@users.noreply.github.com&gt;
Co-authored-by: Ssofja &lt;Ssofja@users.noreply.github.com&gt;
diff --git a/nemo/collections/asr/models/confidence_ensemble.py b/nemo/collections/asr/models/confidence_ensemble.py
@@ -19,7 +19,10 @@
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
-import joblib
+try:
+    from joblib.numpy_pickle_utils import _read_fileobject as _validate_joblib_file
+except ImportError:
+    from joblib.numpy_pickle_utils import _validate_fileobject_and_memmap as _validate_joblib_file
 import numpy as np
 import torch
 from lightning.pytorch import Trainer
@@ -205,13 +208,16 @@ def find_class(self, module, name):
             warnings.simplefilter("ignore")
             # First try to load with our custom unpickler
             try:
-                with open(file_path, 'rb') as f:
-                    unpickler = RestrictedUnpickler(f)
-                    model = unpickler.load()
-            except (pickle.UnpicklingError, AttributeError):
-                # If that fails, try loading with joblib's default loader first
-                # then validate the loaded object
-                model = joblib.load(file_path)
+                with open(file_path, 'rb') as rawf:
+                    with _validate_joblib_file(rawf, file_path, mmap_mode=None) as stream:
+                        if isinstance(stream, tuple):
+                            stream = stream[0]
+
+                        if isinstance(stream, str):
+                            with open(stream, "rb") as f:
+                                model = RestrictedUnpickler(f).load()
+                        else:
+                            model = RestrictedUnpickler(stream).load()
 
                 # Validate the loaded object is a sklearn Pipeline
                 if not isinstance(model, Pipeline):
@@ -222,6 +228,9 @@ def find_class(self, module, name):
                     if not (isinstance(step_obj, (StandardScaler, LogisticRegression))):
                         raise ValueError(f"Unauthorized pipeline step: {type(step_obj)}")
 
+            except (pickle.UnpicklingError, AttributeError) as e:
+                raise SecurityError(f"Failed to safely load model: {e}")
+
         return model
 
     except Exception as e:
