Checksum code now activated by --check flag.

This commit is contained in:
Nikita Skakun 2023-03-30 20:06:12 -07:00
parent d550c12a3e
commit 0cc89e7755

View File

@ -26,6 +26,7 @@ parser.add_argument('--threads', type=int, default=1, help='Number of files to d
parser.add_argument('--text-only', action='store_true', help='Only download text files (txt/json).') parser.add_argument('--text-only', action='store_true', help='Only download text files (txt/json).')
parser.add_argument('--output', type=str, default=None, help='The folder where the model should be saved.') parser.add_argument('--output', type=str, default=None, help='The folder where the model should be saved.')
parser.add_argument('--clean', action='store_true', help='Does not resume the previous download.') parser.add_argument('--clean', action='store_true', help='Does not resume the previous download.')
parser.add_argument('--check', action='store_true', help='Validates the checksums of model files.')
args = parser.parse_args() args = parser.parse_args()
def get_file(url, output_folder): def get_file(url, output_folder):
@ -215,17 +216,18 @@ if __name__ == '__main__':
print(f"Downloading the model to {output_folder}") print(f"Downloading the model to {output_folder}")
download_files(links, output_folder, args.threads) download_files(links, output_folder, args.threads)
# Validate the checksums if args.check:
validated = True # Validate the checksums
for i in range(len(sha256)): validated = True
with open(output_folder / sha256[i][0], "rb") as f: for i in range(len(sha256)):
bytes = f.read() with open(output_folder / sha256[i][0], "rb") as f:
file_hash = hashlib.sha256(bytes).hexdigest() bytes = f.read()
if file_hash != sha256[i][1]: file_hash = hashlib.sha256(bytes).hexdigest()
print(f'[!] Checksum for {sha256[i][0]} failed!') if file_hash != sha256[i][1]:
validated = False print(f'[!] Checksum for {sha256[i][0]} failed!')
validated = False
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('[-] Rerun the download-model.py with --clean flag')