From 5b9b696136f5d6e3a9223983466740785d31b4cc Mon Sep 17 00:00:00 2001 From: Stephen Williamson Date: Sun, 20 Oct 2024 20:46:43 -0400 Subject: [PATCH] Adding byte support for _check_content. --- pickle_scanning_benchmark/download.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pickle_scanning_benchmark/download.py b/pickle_scanning_benchmark/download.py index d1b2b0c..e2067ce 100644 --- a/pickle_scanning_benchmark/download.py +++ b/pickle_scanning_benchmark/download.py @@ -97,8 +97,12 @@ def hf_download_pickle_files( def _check_content(content): # Typical access error from HF API - if content.startswith("Access to model"): - raise Exception("Can not access model file") + if isinstance(content, bytes): + if content.startswith(b"Access to model"): + raise Exception("Can not access model file") + else: + if content.startswith("Access to model"): + raise Exception("Can not access model file") def _download_pickle_file(url, file, outdir):