Add safetensors support

This commit is contained in:
oobabooga 2023-02-03 18:36:32 -03:00
parent 93b0d1b1b8
commit 03f084f311
2 changed files with 4 additions and 1 deletions

View File

@ -71,7 +71,9 @@ if __name__ == '__main__':
for link in links: for link in links:
href = link.get('href')[1:] href = link.get('href')[1:]
if href.startswith(f'{model}/resolve/{branch}'): 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}') downloads.append(f'https://huggingface.co/{href}')
# Downloading the files # Downloading the files

View File

@ -4,3 +4,4 @@ deepspeed==0.8.0
gradio==3.15.0 gradio==3.15.0
transformers==4.25.1 transformers==4.25.1
beautifulsoup4 beautifulsoup4
safetensors