CA-Foundation / backend /download_model.py
β€œvinit5112”
Add all code
deb090d
raw
history blame
2.14 kB
#!/usr/bin/env python3
"""
Download the sentence transformer model for offline use.
Run this script when you have internet access to cache the model locally.
"""
import os
import sys
from sentence_transformers import SentenceTransformer
def download_model():
"""Download and cache the sentence transformer model."""
try:
print("Downloading sentence transformer model 'all-MiniLM-L6-v2'...")
print("This may take a few minutes on first run...")
# This will download and cache the model
model = SentenceTransformer("all-MiniLM-L6-v2")
# Test that it works
test_text = "This is a test sentence."
embedding = model.encode([test_text])
print(f"βœ“ Model downloaded successfully!")
print(f"βœ“ Model tested successfully!")
print(f"βœ“ Embedding dimension: {len(embedding[0])}")
print(f"βœ“ Model cache location: {model.cache_folder}")
return True
except Exception as e:
print(f"βœ— Failed to download model: {e}")
return False
def check_model_exists():
"""Check if the model is already cached."""
try:
# Try to load from cache
import os
os.environ['TRANSFORMERS_OFFLINE'] = '1'
os.environ['HF_HUB_OFFLINE'] = '1'
model = SentenceTransformer("all-MiniLM-L6-v2")
print("βœ“ Model is already cached and available for offline use!")
return True
except Exception:
print("βœ— Model is not cached or not available for offline use")
return False
if __name__ == "__main__":
print("Sentence Transformer Model Downloader")
print("=" * 40)
# Check if model already exists
if check_model_exists():
print("\nModel is already available. No download needed.")
sys.exit(0)
# Download the model
print("\nDownloading model...")
if download_model():
print("\nβœ“ Setup complete! You can now run the application offline.")
else:
print("\nβœ— Download failed. Please check your internet connection.")
sys.exit(1)