Adding byte support for _check_content.

trailofbits · Oct 21, 2024 · 5b9b696 · 5b9b696
1 parent a702d58
commit 5b9b696
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):