Spaces:
Sleeping
Sleeping
File size: 3,744 Bytes
aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 8750f58 aad6058 | 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 | # Model loading utilities - Deployment Optimized
import streamlit as st
import whisper
import easyocr
import os
from sentence_transformers import SentenceTransformer
from transformers import CLIPProcessor, CLIPModel
from gpt4all import GPT4All
from config import MODELS_CONFIG, DEPLOYMENT_MODE
from device_setup import setup_device
@st.cache_resource
def load_models():
try:
# Writable caches first (prevent /.cache writes)
import pathlib
cache_root = os.environ.get("HF_HOME", "/tmp/hf")
pathlib.Path(cache_root).mkdir(parents=True, exist_ok=True)
# Make sure all relevant env vars point to the same place
os.environ.setdefault("HF_HOME", cache_root)
os.environ.setdefault("TRANSFORMERS_CACHE", cache_root)
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", cache_root)
os.environ.setdefault("HF_DATASETS_CACHE", cache_root)
os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", cache_root)
os.environ.setdefault("XDG_CACHE_HOME", "/tmp")
# Setup device
device, has_gpu = setup_device()
device_str = "cuda" if has_gpu else "cpu"
progress = st.progress(0, text="Initializing models...")
# Text embeddings (explicit cache folder)
progress.progress(0.1, text="Loading text encoder...")
st.write("📝 Loading text encoder...")
text_encoder = SentenceTransformer(
MODELS_CONFIG["text_encoder"],
device=device,
cache_folder=cache_root
)
# CLIP (route Hugging Face cache explicitly)
progress.progress(0.3, text="Loading CLIP model...")
st.write("🖼️ Loading CLIP model...")
clip_model_name = MODELS_CONFIG["clip_model"]
clip_model = CLIPModel.from_pretrained(clip_model_name, cache_dir=cache_root).to(device)
clip_processor = CLIPProcessor.from_pretrained(clip_model_name, cache_dir=cache_root)
# Whisper (force CPU on free tier unless GPU present)
progress.progress(0.5, text="Loading Whisper...")
st.write("🎤 Loading Whisper model...")
whisper_size = MODELS_CONFIG["whisper_base"]
whisper_model = whisper.load_model(whisper_size, device=device_str)
# EasyOCR (avoid CUDA probing on CPU Spaces)
progress.progress(0.7, text="Loading OCR...")
st.write("🔍 Loading OCR reader...")
ocr_reader = easyocr.Reader(['en'], gpu=False if not has_gpu else True)
# Local LLM
progress.progress(0.9, text="Loading LLM...")
st.write("🧠 Loading local LLM...")
try:
llm = GPT4All(MODELS_CONFIG["llm_model"], allow_download=True)
st.write("✅ LLM loaded successfully!")
except Exception as llm_error:
st.warning(f"⚠️ LLM loading failed: {llm_error}")
st.info("💡 The system will work without local LLM for document search")
llm = None
progress.progress(1.0, text="All models loaded!")
st.success("🎉 All models loaded successfully!")
return text_encoder, clip_model, clip_processor, whisper_model, ocr_reader, llm, device, has_gpu
except Exception as e:
st.error(f"❌ Error loading models: {e}")
st.info("💡 Try refreshing the page or check your internet connection")
return None, None, None, None, None, None, None, False
def get_model_info():
'''Get information about loaded models'''
return {
"text_encoder": MODELS_CONFIG["text_encoder"],
"clip_model": MODELS_CONFIG["clip_model"],
"whisper_model": MODELS_CONFIG["whisper_base"],
"llm_model": MODELS_CONFIG["llm_model"],
"deployment_mode": DEPLOYMENT_MODE
}
|