diff --git a/tools/fortune_ai/download_model.py b/tools/fortune_ai/download_model.py new file mode 100644 index 0000000..427f495 --- /dev/null +++ b/tools/fortune_ai/download_model.py @@ -0,0 +1,95 @@ +# -*- coding: utf-8 -*- +import hashlib +import json +import os +import sys +import time +import urllib.request + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')) +MANIFEST_PATH = os.path.join(ROOT, 'ai/model.json') + +with open(MANIFEST_PATH, 'r', encoding='utf-8') as f: + manifest = json.load(f) + +dest_path = os.path.join(ROOT, 'ai', manifest['file']) +part_path = dest_path + '.part' +os.makedirs(os.path.dirname(dest_path), exist_ok=True) + +expected_bytes = manifest['bytes'] +expected_sha256 = manifest['sha256'].lower() + +if os.path.exists(dest_path): + actual_bytes = os.path.getsize(dest_path) + if actual_bytes == expected_bytes: + print(f"[INFO] Destination file already exists ({actual_bytes} bytes). Verifying SHA256...") + h = hashlib.sha256() + with open(dest_path, 'rb') as f: + while chunk := f.read(1024 * 1024): + h.update(chunk) + if h.hexdigest().lower() == expected_sha256: + print("[SUCCESS] Model already downloaded and verified!") + sys.exit(0) + else: + print("[WARN] Existing file has checksum mismatch, re-downloading...") + os.remove(dest_path) + +url = f"https://huggingface.co/{manifest['source']}/resolve/{manifest['revision']}/{manifest['source_file']}?download=true" +print(f"Downloading model from:\n{url}") +print(f"Destination: {dest_path}") +print(f"Expected size: {expected_bytes} bytes (~{expected_bytes / (1024**3):.2f} GB)") + +headers = {'User-Agent': 'Mozilla/5.0'} +existing_bytes = 0 +if os.path.exists(part_path): + existing_bytes = os.path.getsize(part_path) + if existing_bytes < expected_bytes: + print(f"Resuming download from byte {existing_bytes}...") + headers['Range'] = f"bytes={existing_bytes}-" + else: + existing_bytes = 0 + os.remove(part_path) + +req = urllib.request.Request(url, headers=headers) +try: + response = urllib.request.urlopen(req, timeout=30) +except Exception as e: + print(f"[ERROR] Failed to open URL: {e}") + sys.exit(1) + +mode = 'ab' if existing_bytes > 0 else 'wb' +downloaded = existing_bytes +start_time = time.time() +last_report_time = start_time +last_report_bytes = downloaded + +with open(part_path, mode) as out: + while True: + chunk = response.read(1024 * 512) # 512 KB + if not chunk: + break + out.write(chunk) + downloaded += len(chunk) + now = time.time() + if now - last_report_time >= 5.0 or downloaded == expected_bytes: + speed_mb = (downloaded - last_report_bytes) / (now - last_report_time) / (1024 * 1024) + pct = (downloaded / expected_bytes) * 100.0 + print(f"Progress: {downloaded / (1024**3):.2f} / {expected_bytes / (1024**3):.2f} GB ({pct:.1f}%) | {speed_mb:.2f} MB/s") + last_report_time = now + last_report_bytes = downloaded + +print("\nDownload finished! Verifying SHA256 checksum...") +h = hashlib.sha256() +with open(part_path, 'rb') as f: + while chunk := f.read(1024 * 1024): + h.update(chunk) + +actual_sha256 = h.hexdigest().lower() +if actual_sha256 != expected_sha256: + print(f"[ERROR] SHA256 mismatch!\nExpected: {expected_sha256}\nGot: {actual_sha256}") + sys.exit(1) + +if os.path.exists(dest_path): + os.remove(dest_path) +os.rename(part_path, dest_path) +print("[SUCCESS] Model successfully downloaded and verified!")