mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add safetensors support

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

4
download-model.py
View file

@ -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

1
requirements.txt
View file

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