From 1aed2b9e52805659b4522cd4b8aaca3677c65a13 Mon Sep 17 00:00:00 2001 From: Morgan Schweers Date: Wed, 31 May 2023 20:11:21 -0700 Subject: [PATCH] Make it possible to download protected HF models from the command line. (#2408) --- README.md | 4 +- download-model.py | 320 ++++++++++++++++++++++++---------------------- server.py | 5 +- 3 files changed, 172 insertions(+), 157 deletions(-) diff --git a/README.md b/README.md index 22134c83..6ac659f0 100644 --- a/README.md +++ b/README.md @@ -156,7 +156,9 @@ For example: python download-model.py facebook/opt-1.3b -If you want to download a model manually, note that all you need are the json, txt, and pytorch\*.bin (or model*.safetensors) files. The remaining files are not necessary. +* If you want to download a model manually, note that all you need are the json, txt, and pytorch\*.bin (or model*.safetensors) files. The remaining files are not necessary. + +* If you want to download a protected model (one gated behind accepting a license or otherwise private, like `bigcode/starcoder`) you can set the environment variables `HF_USER` to your huggingface username and `HF_PASS` to your password or (_as a better option_) to a [User Access Token](https://huggingface.co/settings/tokens). Note that you will need to accept the model terms on the Hugging Face website before starting the download. #### GGML models diff --git a/download-model.py b/download-model.py index 8577f9a1..83eab84a 100644 --- a/download-model.py +++ b/download-model.py @@ -12,6 +12,7 @@ import datetime import hashlib import json import re +import os import sys from pathlib import Path @@ -70,173 +71,183 @@ EleutherAI/pythia-1.4b-deduped return model, branch -def sanitize_model_and_branch_names(model, branch): - if model[-1] == '/': - model = model[:-1] - if branch is None: - branch = "main" - else: - pattern = re.compile(r"^[a-zA-Z0-9._-]+$") - if not pattern.match(branch): - raise ValueError("Invalid branch name. Only alphanumeric characters, period, underscore and dash are allowed.") - - return model, branch +class ModelDownloader: + def __init__(self): + self.s = requests.Session() + if os.getenv('HF_USER') is not None and os.getenv('HF_PASS') is not None: + self.s.auth = (os.getenv('HF_USER'), os.getenv('HF_PASS')) -def get_download_links_from_huggingface(model, branch, text_only=False): - base = "https://huggingface.co" - page = f"/api/models/{model}/tree/{branch}" - cursor = b"" + def sanitize_model_and_branch_names(self, model, branch): + if model[-1] == '/': + model = model[:-1] - links = [] - sha256 = [] - classifications = [] - has_pytorch = False - has_pt = False - has_ggml = False - has_safetensors = False - is_lora = False - while True: - url = f"{base}{page}" + (f"?cursor={cursor.decode()}" if cursor else "") - r = requests.get(url, timeout=10) - r.raise_for_status() - content = r.content + if branch is None: + branch = "main" + else: + pattern = re.compile(r"^[a-zA-Z0-9._-]+$") + if not pattern.match(branch): + raise ValueError( + "Invalid branch name. Only alphanumeric characters, period, underscore and dash are allowed.") - dict = json.loads(content) - if len(dict) == 0: - break - - for i in range(len(dict)): - fname = dict[i]['path'] - if not is_lora and fname.endswith(('adapter_config.json', 'adapter_model.bin')): - is_lora = True - - is_pytorch = re.match("(pytorch|adapter|gptq)_model.*\.bin", fname) - is_safetensors = re.match(".*\.safetensors", fname) - is_pt = re.match(".*\.pt", fname) - is_ggml = re.match(".*ggml.*\.bin", fname) - is_tokenizer = re.match("(tokenizer|ice).*\.model", fname) - is_text = re.match(".*\.(txt|json|py|md)", fname) or is_tokenizer - - if any((is_pytorch, is_safetensors, is_pt, is_ggml, is_tokenizer, is_text)): - if 'lfs' in dict[i]: - sha256.append([fname, dict[i]['lfs']['oid']]) - if is_text: - links.append(f"https://huggingface.co/{model}/resolve/{branch}/{fname}") - classifications.append('text') - continue - if not text_only: - links.append(f"https://huggingface.co/{model}/resolve/{branch}/{fname}") - if is_safetensors: - has_safetensors = True - classifications.append('safetensors') - elif is_pytorch: - has_pytorch = True - classifications.append('pytorch') - elif is_pt: - has_pt = True - classifications.append('pt') - elif is_ggml: - has_ggml = True - classifications.append('ggml') - - cursor = base64.b64encode(f'{{"file_name":"{dict[-1]["path"]}"}}'.encode()) + b':50' - cursor = base64.b64encode(cursor) - cursor = cursor.replace(b'=', b'%3D') - - # If both pytorch and safetensors are available, download safetensors only - if (has_pytorch or has_pt) and has_safetensors: - for i in range(len(classifications) - 1, -1, -1): - if classifications[i] in ['pytorch', 'pt']: - links.pop(i) - - return links, sha256, is_lora + return model, branch -def get_output_folder(model, branch, is_lora, base_folder=None): - if base_folder is None: - base_folder = 'models' if not is_lora else 'loras' + def get_download_links_from_huggingface(self, model, branch, text_only=False): + base = "https://huggingface.co" + page = f"/api/models/{model}/tree/{branch}" + cursor = b"" - output_folder = f"{'_'.join(model.split('/')[-2:])}" - if branch != 'main': - output_folder += f'_{branch}' - output_folder = Path(base_folder) / output_folder - return output_folder + links = [] + sha256 = [] + classifications = [] + has_pytorch = False + has_pt = False + has_ggml = False + has_safetensors = False + is_lora = False + while True: + url = f"{base}{page}" + (f"?cursor={cursor.decode()}" if cursor else "") + r = self.s.get(url, timeout=10) + r.raise_for_status() + content = r.content + + dict = json.loads(content) + if len(dict) == 0: + break + + for i in range(len(dict)): + fname = dict[i]['path'] + if not is_lora and fname.endswith(('adapter_config.json', 'adapter_model.bin')): + is_lora = True + + is_pytorch = re.match("(pytorch|adapter)_model.*\.bin", fname) + is_safetensors = re.match(".*\.safetensors", fname) + is_pt = re.match(".*\.pt", fname) + is_ggml = re.match(".*ggml.*\.bin", fname) + is_tokenizer = re.match("(tokenizer|ice).*\.model", fname) + is_text = re.match(".*\.(txt|json|py|md)", fname) or is_tokenizer + if any((is_pytorch, is_safetensors, is_pt, is_ggml, is_tokenizer, is_text)): + if 'lfs' in dict[i]: + sha256.append([fname, dict[i]['lfs']['oid']]) + + if is_text: + links.append(f"https://huggingface.co/{model}/resolve/{branch}/{fname}") + classifications.append('text') + continue + + if not text_only: + links.append(f"https://huggingface.co/{model}/resolve/{branch}/{fname}") + if is_safetensors: + has_safetensors = True + classifications.append('safetensors') + elif is_pytorch: + has_pytorch = True + classifications.append('pytorch') + elif is_pt: + has_pt = True + classifications.append('pt') + elif is_ggml: + has_ggml = True + classifications.append('ggml') + + cursor = base64.b64encode(f'{{"file_name":"{dict[-1]["path"]}"}}'.encode()) + b':50' + cursor = base64.b64encode(cursor) + cursor = cursor.replace(b'=', b'%3D') + + # If both pytorch and safetensors are available, download safetensors only + if (has_pytorch or has_pt) and has_safetensors: + for i in range(len(classifications) - 1, -1, -1): + if classifications[i] in ['pytorch', 'pt']: + links.pop(i) + + return links, sha256, is_lora -def get_single_file(url, output_folder, start_from_scratch=False): - filename = Path(url.rsplit('/', 1)[1]) - output_path = output_folder / filename - if output_path.exists() and not start_from_scratch: - # Check if the file has already been downloaded completely - r = requests.get(url, stream=True, timeout=10) - total_size = int(r.headers.get('content-length', 0)) - if output_path.stat().st_size >= total_size: - return - # Otherwise, resume the download from where it left off - headers = {'Range': f'bytes={output_path.stat().st_size}-'} - mode = 'ab' - else: - headers = {} - mode = 'wb' + def get_output_folder(self, model, branch, is_lora, base_folder=None): + if base_folder is None: + base_folder = 'models' if not is_lora else 'loras' - r = requests.get(url, stream=True, headers=headers, timeout=10) - with open(output_path, mode) as f: - total_size = int(r.headers.get('content-length', 0)) - block_size = 1024 - with tqdm.tqdm(total=total_size, unit='iB', unit_scale=True, bar_format='{l_bar}{bar}| {n_fmt:6}/{total_fmt:6} {rate_fmt:6}') as t: - for data in r.iter_content(block_size): - t.update(len(data)) - f.write(data) + output_folder = f"{'_'.join(model.split('/')[-2:])}" + if branch != 'main': + output_folder += f'_{branch}' + output_folder = Path(base_folder) / output_folder + return output_folder -def start_download_threads(file_list, output_folder, start_from_scratch=False, threads=1): - thread_map(lambda url: get_single_file(url, output_folder, start_from_scratch=start_from_scratch), file_list, max_workers=threads, disable=True) + def get_single_file(self, url, output_folder, start_from_scratch=False): + filename = Path(url.rsplit('/', 1)[1]) + output_path = output_folder / filename + if output_path.exists() and not start_from_scratch: + # Check if the file has already been downloaded completely + r = self.s.get(url, stream=True, timeout=10) + total_size = int(r.headers.get('content-length', 0)) + if output_path.stat().st_size >= total_size: + return + # Otherwise, resume the download from where it left off + headers = {'Range': f'bytes={output_path.stat().st_size}-'} + mode = 'ab' + else: + headers = {} + mode = 'wb' + + r = self.s.get(url, stream=True, headers=headers, timeout=10) + with open(output_path, mode) as f: + total_size = int(r.headers.get('content-length', 0)) + block_size = 1024 + with tqdm.tqdm(total=total_size, unit='iB', unit_scale=True, bar_format='{l_bar}{bar}| {n_fmt:6}/{total_fmt:6} {rate_fmt:6}') as t: + for data in r.iter_content(block_size): + t.update(len(data)) + f.write(data) -def download_model_files(model, branch, links, sha256, output_folder, start_from_scratch=False, threads=1): - # Creating the folder and writing the metadata - if not output_folder.exists(): - output_folder.mkdir(parents=True, exist_ok=True) - with open(output_folder / 'huggingface-metadata.txt', 'w') as f: - f.write(f'url: https://huggingface.co/{model}\n') - f.write(f'branch: {branch}\n') - f.write(f'download date: {str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))}\n') - sha256_str = '' + def start_download_threads(self, file_list, output_folder, start_from_scratch=False, threads=1): + thread_map(lambda url: self.get_single_file(url, output_folder, start_from_scratch=start_from_scratch), file_list, max_workers=threads, disable=True) + + + def download_model_files(self, model, branch, links, sha256, output_folder, start_from_scratch=False, threads=1): + # Creating the folder and writing the metadata + if not output_folder.exists(): + output_folder.mkdir(parents=True, exist_ok=True) + with open(output_folder / 'huggingface-metadata.txt', 'w') as f: + f.write(f'url: https://huggingface.co/{model}\n') + f.write(f'branch: {branch}\n') + f.write(f'download date: {str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))}\n') + sha256_str = '' + for i in range(len(sha256)): + sha256_str += f' {sha256[i][1]} {sha256[i][0]}\n' + if sha256_str != '': + f.write(f'sha256sum:\n{sha256_str}') + + # Downloading the files + print(f"Downloading the model to {output_folder}") + self.start_download_threads(links, output_folder, start_from_scratch=start_from_scratch, threads=threads) + + + def check_model_files(self, model, branch, links, sha256, output_folder): + # Validate the checksums + validated = True for i in range(len(sha256)): - sha256_str += f' {sha256[i][1]} {sha256[i][0]}\n' - if sha256_str != '': - f.write(f'sha256sum:\n{sha256_str}') + fpath = (output_folder / sha256[i][0]) - # Downloading the files - print(f"Downloading the model to {output_folder}") - start_download_threads(links, output_folder, start_from_scratch=start_from_scratch, threads=threads) - - -def check_model_files(model, branch, links, sha256, output_folder): - # Validate the checksums - validated = True - 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: - bytes = f.read() - file_hash = hashlib.sha256(bytes).hexdigest() - if file_hash != sha256[i][1]: - print(f'Checksum failed: {sha256[i][0]} {sha256[i][1]}') + if not fpath.exists(): + print(f"The following file is missing: {fpath}") validated = False - else: - print(f'Checksum validated: {sha256[i][0]} {sha256[i][1]}') + continue - if validated: - print('[+] Validated checksums of all model files!') - else: - print('[-] Invalid checksums. Rerun download-model.py with the --clean flag.') + with open(output_folder / sha256[i][0], "rb") as f: + bytes = f.read() + file_hash = hashlib.sha256(bytes).hexdigest() + if file_hash != sha256[i][1]: + print(f'Checksum failed: {sha256[i][0]} {sha256[i][1]}') + validated = False + else: + print(f'Checksum validated: {sha256[i][0]} {sha256[i][1]}') + + if validated: + print('[+] Validated checksums of all model files!') + else: + print('[-] Invalid checksums. Rerun download-model.py with the --clean flag.') if __name__ == '__main__': @@ -256,22 +267,23 @@ if __name__ == '__main__': if model is None: model, branch = select_model_from_default_options() + downloader = ModelDownloader() # Cleaning up the model/branch names try: - model, branch = sanitize_model_and_branch_names(model, branch) + model, branch = downloader.sanitize_model_and_branch_names(model, branch) except ValueError as err_branch: print(f"Error: {err_branch}") sys.exit() # Getting the download links from Hugging Face - links, sha256, is_lora = get_download_links_from_huggingface(model, branch, text_only=args.text_only) + links, sha256, is_lora = downloader.get_download_links_from_huggingface(model, branch, text_only=args.text_only) # Getting the output folder - output_folder = get_output_folder(model, branch, is_lora, base_folder=args.output) + output_folder = downloader.get_output_folder(model, branch, is_lora, base_folder=args.output) if args.check: # Check previously downloaded files - check_model_files(model, branch, links, sha256, output_folder) + downloader.check_model_files(model, branch, links, sha256, output_folder) else: # Download files - download_model_files(model, branch, links, sha256, output_folder, threads=args.threads) + downloader.download_model_files(model, branch, links, sha256, output_folder, threads=args.threads) diff --git a/server.py b/server.py index ea8348a9..91c0a1dd 100644 --- a/server.py +++ b/server.py @@ -184,7 +184,8 @@ def count_tokens(text): def download_model_wrapper(repo_id): try: - downloader = importlib.import_module("download-model") + downloader_module = importlib.import_module("download-model") + downloader = downloader_module.ModelDownloader() repo_id_parts = repo_id.split(":") model = repo_id_parts[0] if len(repo_id_parts) > 0 else repo_id branch = repo_id_parts[1] if len(repo_id_parts) > 1 else "main" @@ -369,7 +370,7 @@ def create_model_menus(): shared.gradio['bf16'] = gr.Checkbox(label="bf16", value=shared.args.bf16) shared.gradio['load_in_8bit'] = gr.Checkbox(label="load-in-8bit", value=shared.args.load_in_8bit) shared.gradio['trust_remote_code'] = gr.Checkbox(label="trust-remote-code", value=shared.args.trust_remote_code, info='Make sure to inspect the .py files inside the model folder before loading it with this option enabled.') - + with gr.Box(): gr.Markdown('Transformers 4-bit') with gr.Row():