File size: 1,741 Bytes
b553fa8 | 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 | """
One-time setup: Download all models before running the app
Run this once: python setup_models.py
"""
import os
from sentence_transformers import SentenceTransformer
from rapidocr_onnxruntime import RapidOCR
print("\n" + "="*70)
print("SETTING UP RAG DOCUMENT ASSISTANT - ONE-TIME SETUP")
print("="*70)
# Set cache directory
cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "sentence_transformers")
os.makedirs(cache_dir, exist_ok=True)
print(f"\n📁 Cache directory: {cache_dir}")
# Download embedding model
print("\n[1/2] 📦 Downloading embedding model: all-MiniLM-L6-v2")
print(" Size: ~80MB (this may take 1-2 minutes)")
print(" Downloading...", end="", flush=True)
try:
model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder=cache_dir)
# Test the model
test_text = "This is a test sentence."
embedding = model.encode(test_text)
print(" DONE! ✅")
print(f" Model loaded successfully")
print(f" Embedding dimensions: {len(embedding)}")
print(f" Location: {cache_dir}")
except Exception as e:
print(f"\n ❌ Error: {e}")
exit(1)
# Initialize OCR engine
print("\n[2/2] 📦 Initializing OCR engine (RapidOCR)")
print(" Downloading ONNX models...", end="", flush=True)
try:
ocr_engine = RapidOCR()
print(" DONE! ✅")
print(" OCR engine ready")
except Exception as e:
print(f"\n ❌ Error: {e}")
exit(1)
print("\n" + "="*70)
print("✅ SETUP COMPLETE! All models are cached and ready.")
print("="*70)
print("\n🚀 You can now run your app:")
print(" python -m streamlit run app.py")
print("\n💡 Next time you upload documents, it will be MUCH faster!")
print(" (Models are now cached locally)\n")
|