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

Fix the download script on windows (#6)

Browse source
This commit is contained in:
oobabooga 2023-01-13 09:05:21 -03:00
parent d3bd6a3093
commit fcda5d7107
1 changed files with 25 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

48
download-model.py
View file

@ -27,28 +27,30 @@ def get_file(args):
f.write(data) f.write(data)
t.close() t.close()
model = argv[1] if __name__ == '__main__':
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 model = argv[1]
page = requests.get(url) if model[-1] == '/':
soup = BeautifulSoup(page.content, 'html.parser') model = model[:-1]
links = soup.find_all('a') url = f'https://huggingface.co/{model}/tree/main'
downloads = [] output_folder = Path("models") / model.split('/')[-1]
for link in links: if not output_folder.exists():
href = link.get('href')[1:] output_folder.mkdir()
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 # Finding the relevant files to download
print(f"Downloading the model to {output_folder}...") page = requests.get(url)
pool = multiprocessing.Pool(processes=4) soup = BeautifulSoup(page.content, 'html.parser')
results = pool.map(get_file, [[downloads[i], output_folder] for i in range(len(downloads))]) links = soup.find_all('a')
pool.close() downloads = []
pool.join() 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()