From 03f084f31140eb474e7f7aded0323391e9fc9caa Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Fri, 3 Feb 2023 18:36:32 -0300
Subject: [PATCH] Add safetensors support

---
 download-model.py | 4 +++-
 requirements.txt  | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/download-model.py b/download-model.py
index a83ed17a..220302d9 100644
--- a/download-model.py
+++ b/download-model.py
@@ -71,7 +71,9 @@ if __name__ == '__main__':
     for link in links:
         href = link.get('href')[1:]
         if href.startswith(f'{model}/resolve/{branch}'):
-            if href.endswith(('.json', '.txt')) or (href.endswith('.bin') and 'pytorch_model' in href):
+            is_pytorch = href.endswith('.bin') and 'pytorch_model' in href
+            is_safetensors = href.endswith('.safetensors') and 'model' in href
+            if href.endswith(('.json', '.txt')) or is_pytorch or is_safetensors:
                 downloads.append(f'https://huggingface.co/{href}')
 
     # Downloading the files
diff --git a/requirements.txt b/requirements.txt
index e953a7cb..fd747a9e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,4 @@ deepspeed==0.8.0
 gradio==3.15.0
 transformers==4.25.1
 beautifulsoup4
+safetensors