Make it possible to download protected HF models from the command line. (#2408)

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
 

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):
+class ModelDownloader:
-    if model[-1] == '/':
+    def __init__(self):
-        model = model[:-1]
+        self.s = requests.Session()
-    if branch is None:
+        if os.getenv('HF_USER') is not None and os.getenv('HF_PASS') is not None:
-        branch = "main"
+            self.s.auth = (os.getenv('HF_USER'), os.getenv('HF_PASS'))
-    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
 
 
-def get_download_links_from_huggingface(model, branch, text_only=False):
+    def sanitize_model_and_branch_names(self, model, branch):
-    base = "https://huggingface.co"
+        if model[-1] == '/':
-    page = f"/api/models/{model}/tree/{branch}"
+            model = model[:-1]
-    cursor = b""
 
-    links = []
+        if branch is None:
-    sha256 = []
+            branch = "main"
-    classifications = []
+        else:
-    has_pytorch = False
+            pattern = re.compile(r"^[a-zA-Z0-9._-]+$")
-    has_pt = False
+            if not pattern.match(branch):
-    has_ggml = False
+                raise ValueError(
-    has_safetensors = False
+                    "Invalid branch name. Only alphanumeric characters, period, underscore and dash are allowed.")
-    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
 
-        dict = json.loads(content)
+        return model, branch
-        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
 
 
-def get_output_folder(model, branch, is_lora, base_folder=None):
+    def get_download_links_from_huggingface(self, model, branch, text_only=False):
-    if base_folder is None:
+        base = "https://huggingface.co"
-        base_folder = 'models' if not is_lora else 'loras'
+        page = f"/api/models/{model}/tree/{branch}"
+        cursor = b""
 
-    output_folder = f"{'_'.join(model.split('/')[-2:])}"
+        links = []
-    if branch != 'main':
+        sha256 = []
-        output_folder += f'_{branch}'
+        classifications = []
-    output_folder = Path(base_folder) / output_folder
+        has_pytorch = False
-    return output_folder
+        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):
+    def get_output_folder(self, model, branch, is_lora, base_folder=None):
-    filename = Path(url.rsplit('/', 1)[1])
+        if base_folder is None:
-    output_path = output_folder / filename
+            base_folder = 'models' if not is_lora else 'loras'
-    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'
 
-    r = requests.get(url, stream=True, headers=headers, timeout=10)
+        output_folder = f"{'_'.join(model.split('/')[-2:])}"
-    with open(output_path, mode) as f:
+        if branch != 'main':
-        total_size = int(r.headers.get('content-length', 0))
+            output_folder += f'_{branch}'
-        block_size = 1024
+        output_folder = Path(base_folder) / output_folder
-        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:
+        return output_folder
-            for data in r.iter_content(block_size):
-                t.update(len(data))
-                f.write(data)
 
 
-def start_download_threads(file_list, output_folder, start_from_scratch=False, threads=1):
+    def get_single_file(self, url, output_folder, start_from_scratch=False):
-    thread_map(lambda url: get_single_file(url, output_folder, start_from_scratch=start_from_scratch), file_list, max_workers=threads, disable=True)
+        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):
+    def start_download_threads(self, file_list, output_folder, start_from_scratch=False, threads=1):
-    # Creating the folder and writing the metadata
+        thread_map(lambda url: self.get_single_file(url, output_folder, start_from_scratch=start_from_scratch), file_list, max_workers=threads, disable=True)
-    if not output_folder.exists():
+
-        output_folder.mkdir(parents=True, exist_ok=True)
+
-    with open(output_folder / 'huggingface-metadata.txt', 'w') as f:
+    def download_model_files(self, model, branch, links, sha256, output_folder, start_from_scratch=False, threads=1):
-        f.write(f'url: https://huggingface.co/{model}\n')
+        # Creating the folder and writing the metadata
-        f.write(f'branch: {branch}\n')
+        if not output_folder.exists():
-        f.write(f'download date: {str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))}\n')
+            output_folder.mkdir(parents=True, exist_ok=True)
-        sha256_str = ''
+        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'
+            fpath = (output_folder / sha256[i][0])
-        if sha256_str != '':
-            f.write(f'sha256sum:\n{sha256_str}')
 
-    # Downloading the files
+            if not fpath.exists():
-    print(f"Downloading the model to {output_folder}")
+                print(f"The following file is missing: {fpath}")
-    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]}')
                 validated = False
-            else:
+                continue
-                print(f'Checksum validated: {sha256[i][0]}  {sha256[i][1]}')
 
-    if validated:
+            with open(output_folder / sha256[i][0], "rb") as f:
-        print('[+] Validated checksums of all model files!')
+                bytes = f.read()
-    else:
+                file_hash = hashlib.sha256(bytes).hexdigest()
-        print('[-] Invalid checksums. Rerun download-model.py with the --clean flag.')
+                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)

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():