Merge pull request #122 from sbwilli3/patch-1

benchmarking - update check_content for byte support
trailofbits · Oct 21, 2024 · 18d4b93 · 18d4b93
2 parents a702d58 + 5b9b696
commit 18d4b93
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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):