Harden legacy converter checkpoint loading

python/ctranslate2/converters/fairseq.py

@@ -113,6 +113,7 @@ def __init__(
         fixed_dictionary: Optional[str] = None,
         no_default_special_tokens: bool = False,
         user_dir: Optional[str] = None,
+        unsafe_deserialization: bool = False,
     ):
         """Initializes the Fairseq converter.
 
@@ -125,6 +126,8 @@ def __init__(
           no_default_special_tokens: Require all special tokens to be provided by the user
             (e.g. encoder end token, decoder start token).
           user_dir: Path to the user directory containing custom extensions.
+          unsafe_deserialization: Allow unsafe pickle deserialization when loading
+            trusted legacy checkpoints.
         """
         self._model_path = model_path
         self._data_dir = data_dir
@@ -133,6 +136,7 @@ def __init__(
         self._target_lang = target_lang
         self._no_default_special_tokens = no_default_special_tokens
         self._user_dir = user_dir
+        self._unsafe_deserialization = unsafe_deserialization
 
     def _load(self):
         import fairseq
@@ -147,7 +151,9 @@ def _load(self):
 
         with torch.no_grad():
             checkpoint = torch.load(
-                self._model_path, map_location=torch.device("cpu"), weights_only=False
+                self._model_path,
+                map_location=torch.device("cpu"),
+                weights_only=not self._unsafe_deserialization,
             )
             args = checkpoint["args"] or checkpoint["cfg"]["model"]
 
@@ -329,6 +335,14 @@ def main():
             "including the decoder start token."
         ),
     )
+    parser.add_argument(
+        "--unsafe_deserialization",
+        action="store_true",
+        help=(
+            "Allow loading legacy checkpoints with unsafe pickle deserialization. "
+            "Only enable this option for trusted checkpoints."
+        ),
+    )
     Converter.declare_arguments(parser)
     args = parser.parse_args()
     converter = FairseqConverter(
@@ -339,6 +353,7 @@ def main():
         fixed_dictionary=args.fixed_dictionary,
         no_default_special_tokens=args.no_default_special_tokens,
         user_dir=args.user_dir,
+        unsafe_deserialization=args.unsafe_deserialization,
     )
     converter.convert_from_args(args)
 

python/ctranslate2/converters/opennmt_py.py

@@ -163,19 +163,24 @@ def get_vocabs(vocab):
 class OpenNMTPyConverter(Converter):
     """Converts models generated by OpenNMT-py."""
 
-    def __init__(self, model_path: str):
+    def __init__(self, model_path: str, unsafe_deserialization: bool = False):
         """Initializes the OpenNMT-py converter.
 
         Arguments:
           model_path: Path to the OpenNMT-py PyTorch model (.pt file).
+          unsafe_deserialization: Allow unsafe pickle deserialization when loading
+            trusted legacy checkpoints.
         """
         self._model_path = model_path
+        self._unsafe_deserialization = unsafe_deserialization
 
     def _load(self):
         import torch
 
         checkpoint = torch.load(
-            self._model_path, map_location="cpu", weights_only=False
+            self._model_path,
+            map_location="cpu",
+            weights_only=not self._unsafe_deserialization,
         )
 
         src_vocabs, tgt_vocabs = get_vocabs(checkpoint["vocab"])
@@ -352,9 +357,19 @@ def main():
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
     parser.add_argument("--model_path", required=True, help="Model path.")
+    parser.add_argument(
+        "--unsafe_deserialization",
+        action="store_true",
+        help=(
+            "Allow loading legacy checkpoints with unsafe pickle deserialization. "
+            "Only enable this option for trusted checkpoints."
+        ),
+    )
     Converter.declare_arguments(parser)
     args = parser.parse_args()
-    OpenNMTPyConverter(args.model_path).convert_from_args(args)
+    OpenNMTPyConverter(
+        args.model_path, unsafe_deserialization=args.unsafe_deserialization
+    ).convert_from_args(args)
 
 
 if __name__ == "__main__":