74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
|
|
# Standard library imports
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import subprocess
|
|
|
|
# Third-party imports
|
|
import yaml
|
|
|
|
# Parse command-line arguments
|
|
parser = argparse.ArgumentParser(description="Download files from Hugging Face Hub with authentication.")
|
|
parser.add_argument("repo_id", help="Repository ID (e.g. username/repo-name)")
|
|
parser.add_argument("--quant", type=str, help="Quantization extension to include (e.g. IQ2_XS, q4_0, etc.)")
|
|
args = parser.parse_args()
|
|
|
|
repo_id = args.repo_id
|
|
quant = args.quant
|
|
|
|
if quant:
|
|
print(f"Quant mode: {quant}")
|
|
else:
|
|
print("Quant mode: disabled")
|
|
|
|
with open('config.yaml', 'r') as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
default_location = config['default_location']
|
|
token = config.get('token') # Optional: store token in config.yaml
|
|
|
|
safe_repo_name = repo_id.replace('/', '_').replace('\\', '_')
|
|
download_location = os.path.join(default_location, safe_repo_name)
|
|
os.makedirs(download_location, exist_ok=True)
|
|
|
|
default_location = config['default_location']
|
|
token = config.get('token') # Optional: store token in config.yaml
|
|
|
|
# Authenticate with Hugging Face CLI
|
|
if token:
|
|
try:
|
|
subprocess.run(['hf', 'auth', 'login', '--token', token], check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
print("\n[ERROR] Authentication failed. Please check your Hugging Face token in config.yaml.\n"
|
|
"If you want to add the token to the git credential helper, set 'add_to_git_credential' to True in your config or use the '--add-to-git-credential' flag.\n"
|
|
"Original error: Invalid user token or authentication failed.\n")
|
|
exit(1)
|
|
else:
|
|
print("No token found in config.yaml. Please add your Hugging Face token.")
|
|
exit(1)
|
|
|
|
|
|
|
|
|
|
# Download README.md and config.json (no --include, just as positional arguments)
|
|
download_cmd_1 = [
|
|
'hf', 'download', '--local-dir', download_location, repo_id, 'README.md', 'config.json'
|
|
]
|
|
print("Download command 1:")
|
|
print(' '.join(download_cmd_1))
|
|
subprocess.run(download_cmd_1, check=True)
|
|
print(f"Downloaded ['README.md', 'config.json'] to {download_location}")
|
|
|
|
|
|
# Second hf download call: use --include for quantized files
|
|
if quant:
|
|
quant_pattern = f"*{quant}.*"
|
|
download_cmd_2 = [
|
|
'hf', 'download', '--local-dir', download_location, repo_id,
|
|
'--include', quant_pattern
|
|
]
|
|
print("Download command 2:")
|
|
print(' '.join(download_cmd_2))
|
|
subprocess.run(download_cmd_2, check=True)
|