File size: 4,408 Bytes
e93bfbd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """
Download pre-trained checkpoints so users can skip training from scratch.
Usage:
python download_model.py # interactive menu
python download_model.py --all # download everything
python download_model.py --fusion # fusion model only
python download_model.py --blip-finetuned # fine-tuned BLIP only
"""
import argparse
import os
import sys
from pathlib import Path
BASE_URL = "https://huggingface.co/your-org/medicalai-lightweight/resolve/main"
CHECKPOINT_DIR = Path("./checkpoints")
BLIP_DIR = Path("./blip-xray-finetuned")
ONNX_DIR = CHECKPOINT_DIR / "onnx_full"
def _ensure_dir(path):
path.mkdir(parents=True, exist_ok=True)
def _download_file(url, dest):
import requests
from rich.console import Console
from rich.progress import Progress, BarColumn, DownloadColumn, TextColumn
console = Console()
console.print(f"[cyan]Downloading {url.split('/')[-1]}...[/cyan]")
resp = requests.get(url, stream=True)
resp.raise_for_status()
total = int(resp.headers.get("content-length", 0))
with Progress(
TextColumn("[cyan] Download[/cyan]"),
BarColumn(),
DownloadColumn(),
transient=True,
) as progress:
task = progress.add_task("", total=total)
with open(dest, "wb") as f:
for chunk in resp.iter_content(chunk_size=8192):
f.write(chunk)
progress.update(task, advance=len(chunk))
console.print(f"[green] Saved to {dest}[/green]")
def download_fusion():
_ensure_dir(CHECKPOINT_DIR)
_ensure_dir(ONNX_DIR)
files = [
("fusion_model.pth", CHECKPOINT_DIR / "fusion_model.pth"),
("fusion_full.onnx", ONNX_DIR / "fusion_full.onnx"),
("labels.json", ONNX_DIR / "labels.json"),
]
for fname, dest in files:
url = f"{BASE_URL}/{fname}"
print(f" Would download: {url} -> {dest}")
print()
print("[yellow]Note: Pre-trained checkpoints are not yet hosted.[/yellow]")
print("[yellow]Train locally with: python training.py --mode prepare-data && python training.py --mode train[/yellow]")
print("[yellow]Then export: python quantization.py --mode export-full[/yellow]")
def download_blip():
_ensure_dir(BLIP_DIR)
files = [
"config.json",
"model.safetensors",
"preprocessor_config.json",
"special_tokens_map.json",
"tokenizer.json",
"tokenizer_config.json",
"vocab.txt",
]
for fname in files:
url = f"{BASE_URL}/blip-xray-finetuned/{fname}"
dest = BLIP_DIR / fname
print(f" Would download: {url} -> {dest}")
print()
print("[yellow]Note: Fine-tuned BLIP is not yet hosted.[/yellow]")
print("[yellow]Train locally with: python xray_training.py --mode train[/yellow]")
def main():
parser = argparse.ArgumentParser(description="Download pre-trained models")
parser.add_argument("--all", action="store_true", help="Download everything")
parser.add_argument("--fusion", action="store_true", help="Download fusion model")
parser.add_argument("--blip-finetuned", action="store_true", help="Download fine-tuned BLIP")
args = parser.parse_args()
if not any([args.all, args.fusion, args.blip_finetuned]):
from rich.console import Console
import questionary
console = Console()
console.print("[bold cyan]Download Pre-trained Models[/bold cyan]")
choice = questionary.select(
"What would you like to download?",
choices=[
"Fusion model (Symptom Check) — ONNX + PyTorch",
"Fine-tuned BLIP (Vision captioning)",
"Everything",
"Cancel",
],
).ask()
if choice == "Cancel":
return
if choice == "Fusion model (Symptom Check) — ONNX + PyTorch":
args.fusion = True
elif choice == "Fine-tuned BLIP (Vision captioning)":
args.blip_finetuned = True
else:
args.all = True
try:
import requests
except ImportError:
print("'requests' is required. Install with: pip install requests")
sys.exit(1)
if args.all or args.fusion:
download_fusion()
if args.all or args.blip_finetuned:
download_blip()
if __name__ == "__main__":
main()
|