text-generation-webui/download-model.py
2023-01-13 09:05:21 -03:00

57 lines
1.7 KiB
Python

'''
Downloads models from Hugging Face to models/model-name.
Example:
python download-model.py facebook/opt-1.3b
'''
import requests
from bs4 import BeautifulSoup
import multiprocessing
import tqdm
from sys import argv
from pathlib import Path
def get_file(args):
url = args[0]
output_folder = args[1]
r = requests.get(url, stream=True)
with open(output_folder / Path(url.split('/')[-1]), 'wb') as f:
total_size = int(r.headers.get('content-length', 0))
block_size = 1024
t = tqdm.tqdm(total=total_size, unit='iB', unit_scale=True)
for data in r.iter_content(block_size):
t.update(len(data))
f.write(data)
t.close()
if __name__ == '__main__':
model = argv[1]
if model[-1] == '/':
model = model[:-1]
url = f'https://huggingface.co/{model}/tree/main'
output_folder = Path("models") / model.split('/')[-1]
if not output_folder.exists():
output_folder.mkdir()
# Finding the relevant files to download
page = requests.get(url)
soup = BeautifulSoup(page.content, 'html.parser')
links = soup.find_all('a')
downloads = []
for link in links:
href = link.get('href')[1:]
if href.startswith(f'{model}/resolve/main'):
if href.endswith(('.json', '.txt')) or (href.endswith('.bin') and 'pytorch_model' in href):
downloads.append(f'https://huggingface.co/{href}')
# Downloading the files
print(f"Downloading the model to {output_folder}...")
pool = multiprocessing.Pool(processes=4)
results = pool.map(get_file, [[downloads[i], output_folder] for i in range(len(downloads))])
pool.close()
pool.join()