Spaces:
Running
Running
File size: 2,548 Bytes
2a31b59 | 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 | #!/usr/bin/env python3
"""
Pre-download all models used in HuggingFace Enabling Sessions
Run this BEFORE your session to cache models locally
Models will save to ~/.cache/huggingface/hub
"""
import os
from transformers import pipeline, AutoTokenizer
from sentence_transformers import SentenceTransformer
import config
print("=" * 60)
print("🤗 HuggingFace Model Pre-Download Script")
print("=" * 60)
# Set HF cache directory (optional, for explicit control)
HF_HOME = os.path.expanduser("~/.cache/huggingface")
os.makedirs(HF_HOME, exist_ok=True)
print(f"\n📁 Cache location: {HF_HOME}")
models_to_download = [
("Sentiment Analysis", config.SENTIMENT_MODEL, "sentiment"),
("NER", config.NER_MODEL, "ner"),
("Question Answering", config.QA_MODEL, "qa"),
("Summarization", config.SUMMARIZATION_MODEL, "summarization"),
("Semantic Similarity", config.EMBEDDINGS_MODEL, "embeddings"),
]
print(f"\n📥 Starting download of {len(models_to_download)} models...\n")
# Download pipelines
for task_name, model_id, task_type in models_to_download[:4]:
try:
print(f"⏳ Downloading {task_name} ({model_id})...")
if task_type == "ner":
pipeline("ner", model=model_id)
elif task_type == "qa":
pipeline("question-answering", model=model_id)
elif task_type == "summarization":
pipeline("summarization", model=model_id)
else:
pipeline("sentiment-analysis", model=model_id)
print(f"✅ {task_name} downloaded successfully\n")
except Exception as e:
print(f"❌ Error downloading {task_name}: {str(e)}\n")
# Download Sentence-BERT
try:
print(f"⏳ Downloading Semantic Similarity ({config.EMBEDDINGS_MODEL})...")
SentenceTransformer(config.EMBEDDINGS_MODEL)
print(f"✅ Semantic Similarity downloaded successfully\n")
except Exception as e:
print(f"❌ Error downloading Semantic Similarity: {str(e)}\n")
# Download tokenizer
try:
print(f"⏳ Downloading Tokenizer ({config.SENTIMENT_MODEL})...")
AutoTokenizer.from_pretrained(config.SENTIMENT_MODEL)
print(f"✅ Tokenizer downloaded successfully\n")
except Exception as e:
print(f"❌ Error downloading Tokenizer: {str(e)}\n")
print("=" * 60)
print("✅ Model pre-download complete!")
print("=" * 60)
print("\n📝 Notes:")
print("- All models are cached in ~/.cache/huggingface/hub")
print("- Models will be used instantly in Spaces demos")
print("- Total size: ~2-3 GB (may take 10-20 minutes)")
print("\n🚀 Ready for your session!")
|