Spaces:
Running
Running
import faiss | |
from sentence_transformers import SentenceTransformer | |
import numpy as np | |
import os | |
# Set up cache directory in a writable location | |
cache_dir = os.path.join(os.getcwd(), ".cache") | |
os.makedirs(cache_dir, exist_ok=True) | |
os.environ['HF_HOME'] = cache_dir | |
os.environ['TRANSFORMERS_CACHE'] = cache_dir | |
# Initialize model as None - will be loaded lazily | |
_model = None | |
def preload_model(): | |
"""Preload the sentence transformer model at startup""" | |
global _model | |
if _model is None: | |
print("Preloading sentence transformer model...") | |
try: | |
_model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder=cache_dir) | |
print("Model preloading completed") | |
except Exception as e: | |
print(f"Error loading model: {e}") | |
# Fallback to a different model if the first one fails | |
try: | |
_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir) | |
print("Fallback model preloading completed") | |
except Exception as e2: | |
print(f"Error loading fallback model: {e2}") | |
raise | |
return _model | |
def get_model(): | |
"""Get the sentence transformer model, loading it lazily if needed""" | |
global _model | |
if _model is None: | |
print("Warning: Model not preloaded, loading now...") | |
return preload_model() | |
return _model | |
def build_faiss_index(chunks): | |
model = get_model() | |
embeddings = model.encode(chunks) | |
dimension = embeddings.shape[1] | |
index = faiss.IndexFlatL2(dimension) | |
index.add(np.array(embeddings)) | |
return index, chunks |