From ac118a5ad0f61b094b4be7a73fc993476731936c Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Thu, 4 Apr 2024 13:16:32 +0000 Subject: [PATCH] Pickle conversion now requires `--trust-remote-code`. --- server/text_generation_server/cli.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/server/text_generation_server/cli.py b/server/text_generation_server/cli.py index a513f5e6..f1048297 100644 --- a/server/text_generation_server/cli.py +++ b/server/text_generation_server/cli.py @@ -249,6 +249,13 @@ def download_weights( local_pt_files = utils.download_weights(pt_filenames, model_id, revision) if auto_convert: + if not trust_remote_code: + raise RuntimeError( + f"Safetensors conversion is disabled without `--trust-remote-code` because " + f"Pickle files are unsafe and can essentially contain remote code execution." + f"Please check the safety checks on the hub and ideally make the conversion in a sandbox or in a space." + ) + logger.warning( f"No safetensors weights found for model {model_id} at revision {revision}. " f"Converting PyTorch weights to safetensors."