sanjeevani-api / scripts /download_models.py
Sivaneshakumar's picture
V1
151ec26
Raw History Blame Contribute Delete
1.74 kB
from pathlib import Path
from huggingface_hub import snapshot_download
BASE_DIR = Path(__file__).resolve().parents[1]
MODEL_DIR = BASE_DIR / "models" / "pretrained"
MODELS = {
"trocr": "microsoft/trocr-base-printed",
"biomedbert":
"microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
"sapbert":
"cambridgeltl/SapBERT-from-PubMedBERT-fulltext",
"intent":
"distilbert/distilbert-base-uncased",
"whisper":
"openai/whisper-base",
"indictrans_en_indic":
"ai4bharat/indictrans2-en-indic-dist-200M",
"indictrans_indic_indic":
"ai4bharat/indictrans2-indic-indic-dist-320M",
}
# Only download files required for PyTorch/Transformers.
# This avoids downloading TensorFlow, Flax and duplicate PyTorch weights.
ALLOW_PATTERNS = [
"*.json",
"*.safetensors",
"*.txt",
"*.model",
"*.vocab",
"*.merges",
"*.tiktoken",
"*.py",
]
for name, repo_id in MODELS.items():
destination = MODEL_DIR / name
print("\n" + "=" * 70)
print(f"Downloading: {name}")
print(f"Repository: {repo_id}")
print(f"Destination: {destination}")
print("=" * 70)
destination.mkdir(
parents=True,
exist_ok=True
)
try:
snapshot_download(
repo_id=repo_id,
local_dir=str(destination),
allow_patterns=ALLOW_PATTERNS,
max_workers=2,
)
print(f"\n✅ Finished: {name}")
except Exception as e:
print(f"\n❌ FAILED: {name}")
print(f"Error: {e}")
print(
"\nThe script will continue with the next model."
)
print("\n" + "=" * 70)
print("MODEL DOWNLOAD PROCESS FINISHED")
print("=" * 70)