Issurance_Agent_Rag / embedder.py
Rivalcoder
Add application file
e15840d
raw
history blame
1.25 kB
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 get_model():
"""Get the sentence transformer model, loading it lazily if needed"""
global _model
if _model is None:
try:
_model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder=cache_dir)
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)
except Exception as e2:
print(f"Error loading fallback model: {e2}")
raise
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