mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-10-30 14:10:14 +01:00
55 lines
1.5 KiB
Python
55 lines
1.5 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 os
|
||
|
|
import tqdm
|
||
|
|
from sys import argv
|
||
|
|
|
||
|
|
def get_file(args):
|
||
|
|
url = args[0]
|
||
|
|
output_folder = args[1]
|
||
|
|
|
||
|
|
r = requests.get(url, stream=True)
|
||
|
|
with open(f"{output_folder}/{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()
|
||
|
|
|
||
|
|
model = argv[1]
|
||
|
|
if model.endswith('/'):
|
||
|
|
model = model[:-1]
|
||
|
|
url = f'https://huggingface.co/{model}/tree/main'
|
||
|
|
output_folder = f"models/{model.split('/')[-1]}"
|
||
|
|
if not os.path.exists(output_folder):
|
||
|
|
os.mkdir(output_folder)
|
||
|
|
|
||
|
|
# 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()
|