Don't download if --check is specified

This commit is contained in:
oobabooga 2023-03-31 01:31:47 -03:00
parent 0cc89e7755
commit 92c7068daf

View File

@ -9,6 +9,7 @@ python download-model.py facebook/opt-1.3b
import argparse import argparse
import base64 import base64
import datetime import datetime
import hashlib
import json import json
import re import re
import sys import sys
@ -17,7 +18,6 @@ from pathlib import Path
import requests import requests
import tqdm import tqdm
from tqdm.contrib.concurrent import thread_map from tqdm.contrib.concurrent import thread_map
import hashlib
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('MODEL', type=str, default=None, nargs='?') parser.add_argument('MODEL', type=str, default=None, nargs='?')
@ -212,22 +212,32 @@ if __name__ == '__main__':
if sha256_str != '': if sha256_str != '':
f.write(f'sha256sum:\n{sha256_str}') f.write(f'sha256sum:\n{sha256_str}')
# Downloading the files
print(f"Downloading the model to {output_folder}")
download_files(links, output_folder, args.threads)
if args.check: if args.check:
# Validate the checksums # Validate the checksums
validated = True validated = True
for i in range(len(sha256)): for i in range(len(sha256)):
fpath = (output_folder / sha256[i][0])
if not fpath.exists():
print(f"The following file is missing: {fpath}")
validated = False
continue
with open(output_folder / sha256[i][0], "rb") as f: with open(output_folder / sha256[i][0], "rb") as f:
bytes = f.read() bytes = f.read()
file_hash = hashlib.sha256(bytes).hexdigest() file_hash = hashlib.sha256(bytes).hexdigest()
if file_hash != sha256[i][1]: if file_hash != sha256[i][1]:
print(f'[!] Checksum for {sha256[i][0]} failed!') print(f'Checksum failed: {sha256[i][0]} {sha256[i][1]}')
validated = False validated = False
else:
print(f'Checksum validated: {sha256[i][0]} {sha256[i][1]}')
if validated: if validated:
print('[+] Validated checksums of all model files!') print('[+] Validated checksums of all model files!')
else: else:
print('[-] Rerun the download-model.py with --clean flag') print('[-] Invalid checksums. Rerun download-model.py with the --clean flag.')
else:
# Downloading the files
print(f"Downloading the model to {output_folder}")
download_files(links, output_folder, args.threads)